From 79508c178b8f050ef2e246fae63ee038a17cf6af Mon Sep 17 00:00:00 2001 From: gauenk Date: Wed, 7 Sep 2022 20:57:07 -0400 Subject: [PATCH 1/2] simple setup --- CMakeLists.txt | 25 ------------- setup.py | 100 ++++++------------------------------------------- 2 files changed, 12 insertions(+), 113 deletions(-) delete mode 100644 CMakeLists.txt diff --git a/CMakeLists.txt b/CMakeLists.txt deleted file mode 100644 index ba09e84..0000000 --- a/CMakeLists.txt +++ /dev/null @@ -1,25 +0,0 @@ -cmake_minimum_required(VERSION 3.12...3.18) -project(kepler_jax LANGUAGES CXX) - -message(STATUS "Using CMake version " ${CMAKE_VERSION}) - -find_package(Python COMPONENTS Interpreter Development REQUIRED) -find_package(pybind11 CONFIG REQUIRED) - -include_directories(${CMAKE_CURRENT_LIST_DIR}/lib) - -# CPU op library -pybind11_add_module(cpu_ops ${CMAKE_CURRENT_LIST_DIR}/lib/cpu_ops.cc) -install(TARGETS cpu_ops DESTINATION kepler_jax) - -if (KEPLER_JAX_CUDA) - enable_language(CUDA) - include_directories(${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) - pybind11_add_module( - gpu_ops - ${CMAKE_CURRENT_LIST_DIR}/lib/kernels.cc.cu - ${CMAKE_CURRENT_LIST_DIR}/lib/gpu_ops.cc) - install(TARGETS gpu_ops DESTINATION kepler_jax) -else() - message(STATUS "Building without CUDA") -endif() diff --git a/setup.py b/setup.py index b7e5f9c..93a701c 100644 --- a/setup.py +++ b/setup.py @@ -1,110 +1,34 @@ #!/usr/bin/env python -import codecs -import os -import subprocess - -from setuptools import Extension, find_packages, setup -from setuptools.command.build_ext import build_ext +import os,codecs +from setuptools import find_packages, setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension +# -- read long description -- HERE = os.path.dirname(os.path.realpath(__file__)) - - def read(*parts): with codecs.open(os.path.join(HERE, *parts), "rb", "utf-8") as f: return f.read() - -# This custom class for building the extensions uses CMake to compile. You -# don't have to use CMake for this task, but I found it to be the easiest when -# compiling ops with GPU support since setuptools doesn't have great CUDA -# support. -class CMakeBuildExt(build_ext): - def build_extensions(self): - # First: configure CMake build - import platform - import sys - import distutils.sysconfig - - import pybind11 - - # Work out the relevant Python paths to pass to CMake, adapted from the - # PyTorch build system - if platform.system() == "Windows": - cmake_python_library = "{}/libs/python{}.lib".format( - distutils.sysconfig.get_config_var("prefix"), - distutils.sysconfig.get_config_var("VERSION"), - ) - if not os.path.exists(cmake_python_library): - cmake_python_library = "{}/libs/python{}.lib".format( - sys.base_prefix, - distutils.sysconfig.get_config_var("VERSION"), - ) - else: - cmake_python_library = "{}/{}".format( - distutils.sysconfig.get_config_var("LIBDIR"), - distutils.sysconfig.get_config_var("INSTSONAME"), - ) - cmake_python_include_dir = distutils.sysconfig.get_python_inc() - - install_dir = os.path.abspath( - os.path.dirname(self.get_ext_fullpath("dummy")) - ) - os.makedirs(install_dir, exist_ok=True) - cmake_args = [ - "-DCMAKE_INSTALL_PREFIX={}".format(install_dir), - "-DPython_EXECUTABLE={}".format(sys.executable), - "-DPython_LIBRARIES={}".format(cmake_python_library), - "-DPython_INCLUDE_DIRS={}".format(cmake_python_include_dir), - "-DCMAKE_BUILD_TYPE={}".format( - "Debug" if self.debug else "Release" - ), - "-DCMAKE_PREFIX_PATH={}".format(pybind11.get_cmake_dir()), - ] - if os.environ.get("KEPLER_JAX_CUDA", "no").lower() == "yes": - cmake_args.append("-DKEPLER_JAX_CUDA=yes") - - os.makedirs(self.build_temp, exist_ok=True) - subprocess.check_call( - ["cmake", HERE] + cmake_args, cwd=self.build_temp - ) - - # Build all the extensions - super().build_extensions() - - # Finally run install - subprocess.check_call( - ["cmake", "--build", ".", "--target", "install"], - cwd=self.build_temp, - ) - - def build_extension(self, ext): - target_name = ext.name.split(".")[-1] - subprocess.check_call( - ["cmake", "--build", ".", "--target", target_name], - cwd=self.build_temp, - ) - - +# -- add source files -- extensions = [ - Extension( + CUDAExtension( "kepler_jax.cpu_ops", - ["src/kepler_jax/src/cpu_ops.cc"], + ["lib/cpu_ops.cc"], ), ] - if os.environ.get("KEPLER_JAX_CUDA", "no").lower() == "yes": extensions.append( - Extension( + CUDAExtension( "kepler_jax.gpu_ops", [ - "src/kepler_jax/src/gpu_ops.cc", - "src/kepler_jax/src/cuda_kernels.cc.cu", + "lib/gpu_ops.cc", + "lib/kernels.cc.cu", ], ) ) - +# -- run python make -- setup( name="kepler_jax", author="Dan Foreman-Mackey", @@ -123,5 +47,5 @@ def build_extension(self, ext): install_requires=["jax", "jaxlib"], extras_require={"test": "pytest"}, ext_modules=extensions, - cmdclass={"build_ext": CMakeBuildExt}, + cmdclass={"build_ext": BuildExtension}, ) From a3aff62a164d271bf7aa97a055329fe839031bfe Mon Sep 17 00:00:00 2001 From: gauenk Date: Thu, 8 Sep 2022 12:32:36 -0400 Subject: [PATCH 2/2] added torch --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index a82478d..737c39c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools>=42", "wheel", "setuptools_scm[toml]>=3.4", "pybind11>=2.6", "cmake"] +requires = ["setuptools>=42", "wheel", "setuptools_scm[toml]>=3.4", "pybind11>=2.6", "cmake", "torch"] build-backend = "setuptools.build_meta" [tool.setuptools_scm]