CRAYON/setup_build.py at main · Electroiscoding/CRAYON · GitHub

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
"""
XERV CRAYON SETUP v5.2.6 - WITH C++ EXTENSIONS
==============================================
Builds native extensions for maximum performance
"""

import os
import sys
import platform
from pathlib import Path
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext

VERSION = "5.2.6"

class CustomBuildExt(build_ext):
    """Custom build extension to handle platform-specific compilation"""

    def build_extension(self, ext):
        try:
            super().build_extension(ext)
        except Exception as e:
            print(f"Warning: Failed to build {ext.name}: {e}")
            print("Falling back to pure Python implementation...")
            # Continue without this extension

def get_extensions():
    """Get list of extensions to build"""
    extensions = []

    # Get source directory
    script_dir = Path(__file__).parent
    c_ext_dir = script_dir / "src" / "crayon" / "c_ext"

    # CPU Extension (always try to build)
    cpu_sources = [
        str(c_ext_dir / "crayon_module.c"),
        str(c_ext_dir / "simd_ops.c"),
        str(c_ext_dir / "cpu_engine.cpp"),
    ]

    cpu_ext = Extension(
        'crayon.c_ext.crayon_cpu',
        sources=cpu_sources,
        include_dirs=[str(c_ext_dir)],
        extra_compile_args=['-O3', '-march=native'] if platform.system() != 'Windows' else ['/O2'],
        language='c++'
    )
    extensions.append(cpu_ext)

    # CUDA Extension (optional)
    if os.environ.get('CUDA_HOME') or shutil.which('nvcc'):
        cuda_sources = [
            str(c_ext_dir / "gpu_engine_cuda.cu"),
        ]

        cuda_ext = Extension(
            'crayon.c_ext.crayon_cuda',
            sources=cuda_sources,
            include_dirs=[str(c_ext_dir)],
            extra_compile_args=['-O3'],
            language='c++'
        )
        extensions.append(cuda_ext)

    return extensions

# Check if we should build extensions
build_extensions = '--no-extensions' not in sys.argv

if build_extensions:
    try:
        extensions = get_extensions()
        print(f"Building {len(extensions)} extension(s)...")
    except Exception as e:
        print(f"Warning: Extension setup failed: {e}")
        print("Falling back to pure Python...")
        extensions = []
else:
    extensions = []
    print("Skipping extension build (using pure Python)")

setup(
    name="xerv-crayon",
    version=VERSION,
    packages=find_packages("src"),
    package_dir={"": "src"},
    python_requires=">=3.8,<3.14",
    install_requires=[
        "numpy>=1.21.0",
    ],
    ext_modules=extensions if build_extensions else [],
    cmdclass={'build_ext': CustomBuildExt} if build_extensions else {},
    package_data={
        "crayon": [
            "resources/dat/vocab_lite.dat",
            "resources/dat/vocab_lite.json",
            "resources/dat/vocab_standard.dat",
            "resources/dat/vocab_standard.json",
            "resources/*.txt",
            "resources/*.csv",
            "c_ext/*.h",
            "c_ext/*.c",
            "c_ext/*.cpp",
            "c_ext/*.cu",
            "c_ext/*.hip",
            "c_ext/*.pyd",
            "c_ext/*.so"
        ]
    },
)