利用swig 编写c/c++扩展到Python的接口2说明（win10编译GPU的cu文件）

最新推荐文章于 2022-10-09 10:04:03 发布

meng_zhi_xiang

最新推荐文章于 2022-10-09 10:04:03 发布

阅读量415

点赞数 1

分类专栏：记录1

本文链接：https://blog.csdn.net/meng_zhi_xiang/article/details/106554187

版权

记录1 专栏收录该内容

13 篇文章 1 订阅

订阅专栏

利用swig,Pytorch 编写cuda代码到python的接口

简要说明：

简要说明：

先前写过关于“利用swig 编写c/c++扩展到Python的接口（数组通过numpy传递）”的博客，个人实现了c\c++代码到python的接口，初步调试暂时没发现什么问题；后来因为写了GPU的代码，需要做GPU到python的扩展，需用到nvcc，直接通过python的模块“from distutils.core import setup, Extension”一直扩展不成功，说是找不到编译器编译‘.cu’文件，网上找各种办法，试了好几天，想简单点，各种郁闷，实在没办法，终于想到好像有人通过pytorch实现过相似的转化，然后不小心去pytorch的官方文档（pytorch-Docs > torch.utils.cpp_extension）逛了下，在先前的基础上试了‘Example’代码，结果竟然可以编译通过，然后自己调试了下自己的工程，可以输出和vs，matlab类似的结果，为了方便其他人少走坑，贴下接口代码和相关内容作为一点启发（个人用的anaconda3，win10）；
这个’-arch=sm_61’里面的参数，具体看个人GPU硬件了，是可以查的，自己是从nvcc flags里面查的，个人显卡GTX 1070，所以写的是sm_61，不清楚填其它的有什么影响。

pytorch的官方文档有下面一段代码(官方就是靠谱，真香)，个人就是根据下面这个才解决这几天没解决的编译问题：

Example

>>> from setuptools import setup
>>> from torch.utils.cpp_extension import BuildExtension, CUDAExtension
>>> setup(
        name='cuda_extension',
        ext_modules=[
            CUDAExtension(
                    name='cuda_extension',
                    sources=['extension.cpp', 'extension_kernel.cu'],
                    extra_compile_args={'cxx': ['-g'],
                                        'nvcc': ['-O2']})
        ],
        cmdclass={
            'build_ext': BuildExtension
        })

个人GPU用的是setup.py文件(前面的定位路径查找是网上找到的，谢谢相关人，如果很清楚路径，其实也可以直接填路径，只不过查找的当然很通用)，主要是这个，还有自己先前博客里面的swig相关的i文件：



#!/usr/bin/env python

"""
setup.py file for SWIG C++/Python example
"""
import os
from os.path import join as pjoin
import numpy as np
from setuptools import setup
# from distutils.extension import Extension
from Cython.Distutils import build_ext
from torch.utils.cpp_extension import CppExtension
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

nvcc_bin = 'nvcc.exe'
lib_dir = 'lib/x64'
def find_in_path(name, path):
    "Find a file in a search path"
    for dir in path.split(os.pathsep):
        binpath = pjoin(dir, name)
        if os.path.exists(binpath):
            return os.path.abspath(binpath)
    return None

def locate_cuda():
    """Locate the CUDA environment on the system

    Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64'
    and values giving the absolute path to each directory.

    Starts by looking for the CUDAHOME env variable. If not found, everything
    is based on finding 'nvcc' in the PATH.
    """

    # first check if the CUDAHOME env variable is in use
    if 'CUDA_PATH' in os.environ:
        home = os.environ['CUDA_PATH']
        print("home = %s\n" % home)
        nvcc = pjoin(home, 'bin', nvcc_bin)
    else:
        # otherwise, search the PATH for NVCC
        default_path = pjoin(os.sep, 'usr', 'local', 'cuda', 'bin')
        nvcc = find_in_path(nvcc_bin, os.environ['PATH'] + os.pathsep + default_path)
        if nvcc is None:
            raise EnvironmentError('The nvcc binary could not be '
                'located in your $PATH. Either add it to your path, or set $CUDA_PATH')
        home = os.path.dirname(os.path.dirname(nvcc))
        print("home = %s, nvcc = %s\n" % (home, nvcc))


    cudaconfig = {'home':home, 'nvcc':nvcc,
                  'include': pjoin(home, 'include'),
                  'lib64': pjoin(home, lib_dir)}
    for k, v in cudaconfig.items():
        if not os.path.exists(v):
            raise EnvironmentError('The CUDA %s path could not be located in %s' % (k, v))

    return cudaconfig



CUDA = locate_cuda()
#print(CUDA)

example_module =CUDAExtension('_pyGpuPSFSimuf',
                sources=[   'pyGpuPSFSimuf.cpp', 
                'pyGpuPSFSimuf_wrap.cxx',
                './src/psf_gpu_test.cu',
                './src/VectorPSFFinal.cu',], 
              library_dirs=[CUDA['lib64']],
              libraries=['cudart','cufft',],
              language='c++',
              # runtime_library_dirs=[CUDA['lib64']],
              # this syntax is specific to this build system
              # we're only going to use certain compiler args with nvcc and not with
              # gcc the implementation of this trick is in customize_compiler() below
              extra_compile_args={'cxx':[],#'gcc': ["-Wno-unused-function"],
                                  'nvcc': ['-arch=sm_61',
                                           '--ptxas-options=-v',
                                           '-c',
                                           '--compiler-options',
                                           ]},
              include_dirs=[np.get_include(), CUDA['include']]
              )

setup (
    name = 'gpuPSF',
    ext_modules = [example_module],
    cmdclass={

    'build_ext':BuildExtension
    }

)