setup.py

import os
import shutil
import subprocess

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

HERE = os.path.dirname(os.path.abspath(__file__))
PKG = os.path.join(HERE, "src", "mpsparse")
KERNELS = os.path.join(PKG, "kernels")
METALLIB_DIR = os.path.join(PKG, "_metallib")
METAL_CPP = os.path.join(HERE, "third_party", "metal-cpp")

# Metal kernels shared by both extensions. Compiled once into the bundled
# _metallib/ directory and shipped as package data, so installs never need
# the Metal toolchain to *load* them (only to build).
KERNEL_NAMES = ["coo_csr", "spmv", "cgd", "bicgstab", "bcsr_spmv"]


def compile_metallibs():
    os.makedirs(METALLIB_DIR, exist_ok=True)
    for name in KERNEL_NAMES:
        src = os.path.join(KERNELS, f"{name}.metal")
        air = os.path.join(METALLIB_DIR, f"{name}.air")
        lib = os.path.join(METALLIB_DIR, f"{name}.metallib")
        print(f"[mpsparse] compiling {name}.metal -> {name}.metallib")
        subprocess.check_call(["xcrun", "-sdk", "macosx", "metal", "-c", src, "-o", air])
        subprocess.check_call(["xcrun", "-sdk", "macosx", "metallib", air, "-o", lib])
        os.remove(air)


compile_metallibs()


def extension(name, source):
    return CppExtension(
        name=name,
        sources=[os.path.join("src", "mpsparse", source)],
        include_dirs=[METAL_CPP],
        extra_compile_args=["-std=c++17"],
        extra_link_args=["-framework", "Metal", "-framework", "Foundation",
                         "-framework", "QuartzCore"],
    )


setup(
    ext_modules=[
        extension("mpsparse._csr", "spmv.cpp"),
        extension("mpsparse._bcsr", "bcsr.cpp"),
    ],
    cmdclass={"build_ext": BuildExtension},
)