In C++/CUDA extension, how do we include the py wrapper file into setup.py as well?

C++
1

Hi,
In the official C++/CUDA extension tutorial (GitHub - pytorch/extension-cpp: C++ extensions in PyTorch), setup.py only installs the c++ part of the extension into python’s site-packages directory, leaving the lltm.py wrapper file in local directory.
For usage, we will have to refer to this lltm.py file locally (e.g. import cuda.lltm as LLTM, where cuda is the local folder name), which is not ideal for package distribution.
Is there a way to include lltm.py into the extension installation as well?

The setup.py for cuda version looks like this:

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name='lltm_cuda',
    ext_modules=[
        CUDAExtension('lltm_cuda', [
            'lltm_cuda.cpp',
            'lltm_cuda_kernel.cu',
        ]),
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

Thanks!

2

After reading the document of setup script (2. Writing the Setup Script — Python 3.10.1 documentation), i modify the setup.py file (for cuda folder) as follows, and we can import lltm as installed module :

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name='lltm_cuda',
    py_modules = ['lltm'],
    ext_modules=[
        CUDAExtension('lltm_cuda', [
            'lltm_cuda.cpp',
            'lltm_cuda_kernel.cu',
        ]),
    ],
    cmdclass={
        'build_ext': BuildExtension
    })