dev_AI_framework

build 과정의 고민, setup.py 파일 빌드 단순화 방법

명징직조지훈 2024. 12. 23. 07:45
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
import subprocess
import os

class BuildExtWithNvcc(build_ext):
    def build_extensions(self):
        for ext in self.extensions:
            if isinstance(ext, CUDAExtension):
                self.build_cuda_extension(ext)
            else:
                super().build_extensions()

    def build_cuda_extension(self, ext):
        sources = [os.path.abspath(src) for src in ext.sources]
        output_file = os.path.abspath(self.get_ext_fullpath(ext.name))
        build_dir = os.path.dirname(output_file)
        
        # Pybind11 및 Python 경로 설정
        pybind_include = r"C:\Users\owner\AppData\Local\Programs\Python\Python312\Lib\site-packages\pybind11\include"
        python_include = r"C:\Users\owner\AppData\Local\Programs\Python\Python312\include"
        python_lib = r"C:\Users\owner\AppData\Local\Programs\Python\Python312\libs"

        # CUDA 경로 설정
        cuda_lib_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.6\lib\x64"
        cuda_bin_path = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.6\bin"

        # 출력 디렉토리 생성
        os.makedirs(build_dir, exist_ok=True)

        # NVCC 명령 설정
        nvcc_cmd = [
            "nvcc",
            "-shared",                      # 공유 라이브러리 생성
            "-O2",                          # 최적화
            "-x", "cu",                     # CUDA 코드
            f'-I"{pybind_include}"',        # Pybind11 헤더 경로
            f'-I"{python_include}"',        # Python 헤더 경로
            f'-L"{python_lib}"',            # Python 라이브러리 경로 추가
            f'-L"{cuda_lib_path}"',         # CUDA 라이브러리 경로 추가
            "--compiler-options", "/MD",    # MSVC 컴파일러 옵션
            "-lcudart",                     # CUDA 런타임 라이브러리
            "-lpython312",                  # Python 라이브러리 연결
            "-o", f'"{output_file}"',       # 출력 파일 설정
        ] + sources

        # 링커 로그 활성화 (디버깅용)
        nvcc_cmd.extend(["--verbose", "-Xlinker", "/VERBOSE"])

        print("Running NVCC:", " ".join(nvcc_cmd))

        # CUDA DLL 경로를 PATH에 추가
        os.environ["PATH"] += os.pathsep + cuda_bin_path

        # NVCC 실행
        subprocess.check_call(" ".join(nvcc_cmd), shell=True)

class CUDAExtension(Extension):
    def __init__(self, name, sources):
        super().__init__(name, sources)
        self.sources = sources

setup(
    name="cuda_add_example",
    version="0.1",
    ext_modules=[
        CUDAExtension("cuda_add", ["cuda_add.cu"])  # CUDA 확장 모듈
    ],
    cmdclass={"build_ext": BuildExtWithNvcc},       # 빌드 확장 클래스 등록
    zip_safe=False,
)

setup.py 의 내용은 프로젝트마다달라진다. 확장 모듈이 달라지는 것과 빌드 환경이 바뀔 때, 해당 구성이 변경된다.

 

1. 확장 모듈 이름 또는 소스 파일의 변경

빌드하려는 확장 모듈의 이름이나 소스 파일 경로가 달라질 경우 setup.py 의 Extension 정의를 수정해야 한다.

ext_modules=[
    CUDAExtension("module_name", ["source1.cu", "source2.cu"])
]

 

위의 예시 외에도 Python 버전, pybind11 경로, CUDA 버전 등의 변경 시 setup 파일의 내용이 변경되어야 한다.

 

setup 파일을 매번 수정하지 않고, 확장 모듈 이름과 소스 파일을 동적으로 지정하거나 재사용 가능한 방법을 구현

 

1. 동적으로 확장 모듈 지정

명령줄 인자를 받아 확장 모듈 이름과 소스 파일을 동적으로 설정할 수 있다.

 

2. 사용 방법

명령줄에서 setup.py 실행 시 모듈 이름과 소스 파일을 인자로 전달한다.

python setup.py build_ext --name operations_matrix_cuda --sources operations_matrix_cuda.cu

--name: 확장 모듈 이름

--sources : CUDA 소스 파일 경로