ct2rs 0.9.21

Rust bindings for OpenNMT/CTranslate2
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import os
import shutil
import subprocess

from setuptools import Command, Extension, setup, find_packages, find_namespace_packages
from setuptools.command.build_py import build_py
from setuptools.command.build_ext import build_ext
from wheel.bdist_wheel import bdist_wheel


project_path = os.path.abspath(os.path.dirname(__file__))
cccl_path = os.path.abspath(os.path.join(project_path, "..", '..'))
cccl_headers = [
    ['cub', 'cub'],
    ['libcudacxx', 'include'],
    ['thrust', 'thrust']
]
with open(os.path.join(project_path, 'cuda', 'parallel', '_version.py')) as f:
    exec(f.read())
ver = __version__
del __version__


with open("README.md") as f:
    long_description = f.read()


class CustomBuildCommand(build_py):
    def run(self):
        self.run_command('package_cccl')
        build_py.run(self)


class CustomWheelBuild(bdist_wheel):

    def run(self):
        self.run_command('package_cccl')
        super().run()


class PackageCCCLCommand(Command):
    description = 'Generate additional files'
    user_options = []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self):
        for proj_dir, header_dir in cccl_headers:
            src_path = os.path.abspath(
                os.path.join(cccl_path, proj_dir, header_dir))
            # TODO Extract cccl headers into a standalone package
            dst_path = os.path.join(project_path, 'cuda', '_include', proj_dir)
            if os.path.exists(dst_path):
                shutil.rmtree(dst_path)
            shutil.copytree(src_path, dst_path)


class CMakeExtension(Extension):
    def __init__(self, name):
        super().__init__(name, sources=[])


class BuildCMakeExtension(build_ext):
    def run(self):
        for ext in self.extensions:
            self.build_extension(ext)

    def build_extension(self, ext):
        extdir = os.path.abspath(os.path.dirname(
            self.get_ext_fullpath(ext.name)))
        cmake_args = [
            '-DCCCL_ENABLE_CUB=YES',
            '-DCCCL_ENABLE_THRUST=YES',
            '-DCCCL_ENABLE_C=YES',
            '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + extdir,
            '-DCMAKE_BUILD_TYPE=Release',
        ]

        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)

        subprocess.check_call(['cmake', cccl_path] +
                              cmake_args, cwd=self.build_temp)
        subprocess.check_call(
            ['cmake', '--build', '.', '--target', 'cccl.c'], cwd=self.build_temp)


setup(
    name="cuda-parallel",
    version=ver,
    description="Experimental Core Library for CUDA Python",
    long_description=long_description,
    long_description_content_type="text/markdown",
    author="NVIDIA Corporation",
    classifiers=[
        "Programming Language :: Python :: 3 :: Only",
        "Environment :: GPU :: NVIDIA CUDA",
    ],
    packages=find_namespace_packages(include=['cuda.*']),
    python_requires='>=3.9',
    install_requires=[
        "numba>=0.60.0",
        "cuda-python",
        "jinja2"
    ],
    extras_require={
        "test": [
            "pytest",
        ]
    },
    cmdclass={
        'package_cccl': PackageCCCLCommand,
        'build_py': CustomBuildCommand,
        'bdist_wheel': CustomWheelBuild,
        'build_ext': BuildCMakeExtension
    },
    ext_modules=[CMakeExtension('cuda.parallel.experimental.cccl.c')],
    include_package_data=True,
    license="Apache-2.0 with LLVM exception",
    license_files=('../../LICENSE',),
)