通过Python Distutils(用于Python C扩展)使用可重定位的设备代码编译CUDA代码

compile cuda code with relocatable device code through python distutils (for python c extension)

本文关键字:Python 代码 定位 编译 CUDA 用于 Distutils 扩展 通过      更新时间:2023-10-16

我有一些使用协作组的cuda代码,因此需要-rdc=true标志来编译nvcc。我想从 python 调用 cuda 代码,所以我正在编写一个带有 python c 扩展的 python 接口。

因为我包含 cuda 代码,所以我不得不调整我的 setup.py,如中所述:python distutils 可以编译 CUDA 代码吗?

这会编译和安装,但是一旦我在 python 中导入我的代码,它就会出现段错误。删除 -rdc=true 标志使一切正常,但迫使我从 cuda 内核中删除任何协作组代码(或在编译过程中收到"cudaCGGetIntrinsicHandle 未解决"错误(。

有什么方法可以进一步调整我的 setup.py 以使其工作吗?或者,有没有其他方法可以编译我的 c 扩展,允许 cuda 代码(打开 rdc 标志(?

我想我有点想出了答案。如果使用 nvcc 生成可重定位的设备代码,则 nvcc 需要链接目标文件,以便正确处理设备代码链接,或者需要通过在具有带有"--device-link"标志的可重定位设备代码的所有对象文件上运行 nvcc 来生成单独的对象文件。然后,可以将此额外的对象文件包含在外部链接器的所有其他对象文件中。

我从 python distutils 可以编译 CUDA 代码吗?中调整了设置,方法是在源文件列表的末尾添加一个虚拟的"link.cu"文件。我还为 cuda 设备链接步骤添加了 cudadevrt 库和另一组编译器选项:

ext = Extension('mypythonextension',
sources=['python_wrapper.cpp', 'file_with_cuda_code.cu', 'link.cu'],
library_dirs=[CUDA['lib64']],
libraries=['cudart', 'cudadevrt'],
runtime_library_dirs=[CUDA['lib64']],
extra_compile_args={'gcc': [],
'nvcc': ['-arch=sm_70', '-rdc=true', '--compiler-options', "'-fPIC'"],
'nvcclink': ['-arch=sm_70', '--device-link', '--compiler-options', "'-fPIC'"]
},
include_dirs = [numpy_include, CUDA['include'], 'src'])

然后,适应编译器调用的函数会以以下方式获取此问题:

def customize_compiler_for_nvcc(self):
self.src_extensions.append('.cu')
# track all the object files generated with cuda device code
self.cuda_object_files = []
super = self._compile
def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
# generate a special object file that will contain linked in
# relocatable device code
if src == 'link.cu':
self.set_executable('compiler_so', CUDA['nvcc'])
postargs = extra_postargs['nvcclink']
cc_args = self.cuda_object_files[1:]
src = self.cuda_object_files[0]
elif os.path.splitext(src)[1] == '.cu':
self.set_executable('compiler_so', CUDA['nvcc'])
postargs = extra_postargs['nvcc']
self.cuda_object_files.append(obj)
else:
postargs = extra_postargs['gcc']
super(obj, src, ext, cc_args, postargs, pp_opts)
self.compiler_so = default_compiler_so
self._compile = _compile

由于我缺乏distutils知识,该解决方案感觉有点黑客,但它似乎有效。 :)