Continuing the discussion from From torch.utils.ffi import _wrap_function replacement:
Thank u for your reply. I already change into cpp extension, but the older codes has some wrapper using this function. It seems the old code can not word anymore.
But I need a wrapper for some cuda extensions, here is my attempt
"""
this file build extensions under src/

"""
import os
import torch
# from torch.utils.ffi import create_extension
from torch.utils.cpp_extension import BuildExtension, CppExtension
from setuptools import setup, find_packages

sources = ['src/lib_cffi.cpp', 'src/bn.cu']
include_dirs = [os.path.realpath('./src')]
print('include dir: ', include_dirs)

extra_objects = ['src/bn.o']
with_cuda = True

this_file = os.path.dirname(os.path.realpath(__file__))
extra_objects = [os.path.join(this_file, fname) for fname in extra_objects]

if torch.cuda.is_available():
    enable_gpu = True
else:
    enable_gpu = False

if enable_gpu:
    from torch.utils.cpp_extension import CUDAExtension
    build_extension = CUDAExtension
else:
    build_extension = CppExtension

# ffi = create_extension(
#     '_ext',
#     headers=headers,
#     sources=sources,
#     relative_to=__file__,
#     with_cuda=with_cuda,
#     extra_objects=extra_objects,
#     extra_compile_args=["-std=c++11"]
# )

# if __name__ == '__main__':
#     ffi.build()

ext_modules = [
    build_extension(
        '_ext',
        include_dirs=include_dirs,
        sources=sources,
        extra_objects=extra_objects,
        extra_compile_args=["-std=c++11"]
    )
]


setup(
    name='ext',
    version='0.1',
    packages=find_packages(),
    ext_modules=ext_modules,
    cmdclass={'build_ext': BuildExtension}
)

I have these files under src:
bn.cu
bn.h
lib_cffi.cpp
lib_cffi.h

How should I write a wrapper to access them? I tried the setup.py, it got nothing to import in python…
Any suggestions? Here is the source locate:

  
      
      GitHub
  
  
    

jinfagang/LightNet 28

LightNet: Light-weight Networks for Semantic Image Segmentation (Cityscapes and Mapillary Vistas Dataset) - jinfagang/LightNet