amd-repro/setup.py at main · divchenko/amd-repro · 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
"""
    Copyright (c) 2024, Fireworks AI.
"""

from setuptools import setup
from torch.utils import cpp_extension

setup(
    ext_modules=[
        cpp_extension.CUDAExtension(
            name="amd_repro.lib",
            sources=[
                "amd_repro/registry.cpp",
                "amd_repro/virtual_memory.cpp",
                "amd_repro/multi_gpu_barrier.hip",
            ],
            include_dirs=[
                "amd_repro",
            ],
            extra_compile_args={
                "cxx": [
                    "-std=c++17",
                    "-O3",
                ],
                "nvcc": [
                    "-std=c++17",
                    "-O3",
                    "-DNDEBUG",
                    "-g",
                    # hip-specific
                    "-save-temps",
                    "-U__HIP_NO_HALF_CONVERSIONS__",
                    "-U__HIP_NO_HALF_OPERATORS__",
                    "-U__HIP_NO_HALF2_OPERATORS__",
                    "-U__HIP_NO_BFLOAT16_CONVERSIONS__",
                ],
            },
        ),
    ],
    cmdclass={
        "build_ext": cpp_extension.BuildExtension,
    },
)