pytorch setup.py阅读记录

LogicOverlord
213 阅读3分钟

注：本文为个人阅读记录之用，因此写起来比较随心所欲，而且由于是阅读的时候的随手记录，因此错误应该也不少。。。。。

近期打算花时间好好学习一下pytorch的源码，而要学习pytorch的源码，自然就要从setup.py开始了，本文先着重于pytorch v1.13.0 7c98e70d44abc7a1aead68b6ea6c8adc8c554db5, docker为nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04，安装方法为python setup.py install。

setup.py中，首先执行的是里面的main方法，下面介绍main方法里面调用的几个比较大的函数：

①build_deps函数，该函数的调用位于setup.py的990行，该函数首先确定submodule都git clone了，然后调用build_caffe2函数来构建caffe2。

1)build_caffe2函数:该函数的调用位于setup.py的401行，/home/pytorch/tools/setup_helpers/cmake.py中通过check_call在/home/pytorch/build目录下执行如下cmake命令：

cmake -DBUILD_PYTHON=True -DBUILD_TEST=True -DCMAKE_BUILD_TYPE=Release -DCMAKE_EXPORT_COMPILE_COMMANDS=1 -DCMAKE_INSTALL_PREFIX=/home/pytorch/torch -DCMAKE_PREFIX_PATH=/usr/lib/python3.8/site-packages -DNUMPY_INCLUDE_DIR=/usr/local/lib/python3.8/dist-packages/numpy/core/include -DPYTHON_EXECUTABLE=/usr/bin/python -DPYTHON_INCLUDE_DIR=/usr/include/python3.8 -DPYTHON_LIBRARY=/usr/lib/libpython3.8.so.1.0 -DTORCH_BUILD_VERSION=1.13.0a0+git7c98e70 -DUSE_NUMPY=True /home/pytorch

ok，接下来再来看一下cmakelist.txt里面的内容。/home/pytorch/CMakeLists.txt中除了设置编译选项一类的“废话”外，最主要的是1034和1035行的两个add_subdirectory:

#/home/pytorch/CMakeLists.txt
add_subdirectory(c10)
add_subdirectory(caffe2)

c10文件夹介绍:c10/, which is a pun on Caffe2 and A"Ten" (get it? Caffe 10) contains the core abstractions of PyTorch, including the actual implementations of the Tensor and Storage data structures.

/home/pytorch/c10/CMakeLists.txt中主要就是这么几行代码比较重要:

file(GLOB C10_SRCS
        *.cpp
        core/*.cpp
        core/impl/*.cpp
        mobile/*.cpp
        macros/*.cpp
        util/*.cpp
        )
file(GLOB C10_HEADERS
        *.h
        core/*.h
        core/impl/*.h
        mobile/*.h
        macros/*.h
        util/*.h
        )
add_library(c10 ${C10_SRCS} ${C10_HEADERS})
if(USE_CUDA)
  add_subdirectory(cuda)
endif()

可见其就是把c10中的源文件编译成动态库。

/home/pytorch/caffe2/CMakeLists.txt则比c10要庞大的多，其过于复杂就随便看看不做记录了。。。。 最后生成的so库如下:

./build/nccl/lib/libnccl.so
./build/lib/libcaffe2_nvrtc.so
./build/lib/libtorch_cpu.so
./build/lib/libc10.so
./build/lib/libtorch_global_deps.so
./build/lib/libnnapi_backend.so
./build/lib/libc10_cuda.so
./build/lib/libtorch_cuda.so
./build/lib/libtorch.so
./build/lib/libtorchbind_test.so
./build/lib/libtorch_python.so
./build/lib/libbackend_with_compiler.so
./build/lib/libc10d_cuda_test.so
./build/lib/libjitbackend_test.so
./build/lib/libshm.so
./build/lib/libtorch_cuda_linalg.so
./build/functorch/functorch.so

②configure_extension_build函数：位于setup.py的992行，其中配置了extension的相关信息，代码如下:

    C = Extension("torch._C",
                  libraries=main_libraries,  #libtorch_python.so
                  sources=main_sources, #torch/csrc/stub.c
                  language='c',
                  extra_compile_args=main_compile_args + extra_compile_args,
                  include_dirs=[],
                  library_dirs=library_dirs,
                  extra_link_args=extra_link_args + main_link_args + make_relative_rpath_args('lib'))
    C_flatbuffer = Extension("torch._C_flatbuffer",
                             libraries=main_libraries,  #libtorch_python.so
                             sources=["torch/csrc/stub_with_flatbuffer.c"],
                             language='c',
                             extra_compile_args=main_compile_args + extra_compile_args,
                             include_dirs=[],
                             library_dirs=library_dirs,
                             extra_link_args=extra_link_args + main_link_args + make_relative_rpath_args('lib'))

最后生成的编译命令如下:

x86_64-linux-gnu-gcc -pthread -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -g -fwrapv -O2 -fPIC -I/usr/include/python3.8 -c torch/csrc/stub.c -o build/temp.linux-x86_64-cpython-38/torch/csrc/stub.o -Wall -Wextra -Wno-strict-overflow -Wno-unused-parameter -Wno-missing-field-initializers -Wno-write-strings -Wno-unknown-pragmas -Wno-deprecated-declarations -fno-strict-aliasing -Wno-missing-braces
x86_64-linux-gnu-gcc -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fwrapv -O2 build/temp.linux-x86_64-cpython-38/torch/csrc/stub.o -L/home/pytorch/torch/lib -L/usr/local/cuda/lib64/stubs -L/usr/lib -ltorch_python -o build/lib.linux-x86_64-cpython-38/torch/_C.cpython-38-x86_64-linux-gnu.so -Wl,-rpath,$ORIGIN/lib
building 'torch._C_flatbuffer' extension
x86_64-linux-gnu-gcc -pthread -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -g -fwrapv -O2 -fPIC -I/usr/include/python3.8 -c torch/csrc/stub_with_flatbuffer.c -o build/temp.linux-x86_64-cpython-38/torch/csrc/stub_with_flatbuffer.o -Wall -Wextra -Wno-strict-overflow -Wno-unused-parameter -Wno-missing-field-initializers -Wno-write-strings -Wno-unknown-pragmas -Wno-deprecated-declarations -fno-strict-aliasing -Wno-missing-braces
x86_64-linux-gnu-gcc -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -Wl,-z,relro -g -fwrapv -O2 build/temp.linux-x86_64-cpython-38/torch/csrc/stub_with_flatbuffer.o -L/home/pytorch/torch/lib -L/usr/local/cuda/lib64/stubs -L/usr/lib -ltorch_python -o build/lib.linux-x86_64-cpython-38/torch/_C_flatbuffer.cpython-38-x86_64-linux-gnu.so -Wl,-rpath,$ORIGIN/lib

然后就只剩下把头文件 so文件copy到安装目录之类的操作了。

avatar
文章
阅读
粉丝