metal-quant-ext2/setup.py at main · jedt/metal-quant-ext2 · 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
from setuptools import setup
from torch.utils import cpp_extension

pytorch_compile_args = [
    "-std=c++17",
    "-ObjC++",
    "-arch", "arm64",
    "-mmacosx-version-min=12.0",
]
pytorch_link_args = [
    "-arch", "arm64",
    "-mmacosx-version-min=12.0",
]

framework_args = [
    '-framework', 'Foundation',
    '-framework', 'Metal',
    '-framework', 'Accelerate',
    '-framework', 'MetalPerformanceShaders',
    '-framework', 'MetalPerformanceShadersGraph'
]

pytorch_link_args.extend(framework_args)

setup(
    name="metal_quant_ext2",
    version="0.0.1",
    description="A Custom Pytorch MPS extension using MLX",
    ext_modules=[
        cpp_extension.CppExtension(
            "metal_quant_ext2",
            ["CustomOP.mm"],
            extra_compile_args=pytorch_compile_args,
            extra_link_args=pytorch_link_args,
        )
    ],
    cmdclass={"build_ext": cpp_extension.BuildExtension},
)