From 4ee722d2bcabb785b06989274831e3a5d0c12071 Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Tue, 12 Mar 2024 16:35:36 -0700 Subject: [PATCH] Add python runtime wheel. --- CMakeLists.txt | 8 + python_projects/runtime/.gitignore | 4 + python_projects/runtime/README.md | 4 + python_projects/runtime/pyproject.toml | 11 ++ python_projects/runtime/setup.py | 145 +++++++++++++++++++ python_projects/runtime/therock/__init__.py | 37 +++++ python_projects/runtime/therock/test_dist.py | 53 +++++++ 7 files changed, 262 insertions(+) create mode 100644 python_projects/runtime/.gitignore create mode 100644 python_projects/runtime/README.md create mode 100644 python_projects/runtime/pyproject.toml create mode 100644 python_projects/runtime/setup.py create mode 100644 python_projects/runtime/therock/__init__.py create mode 100644 python_projects/runtime/therock/test_dist.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 78d6f32..e41983e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -54,6 +54,7 @@ endmacro() set(DEFAULT_CMAKE_ARGS -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} + -DCMAKE_PLATFORM_NO_VERSIONED_SONAME=${CMAKE_PLATFORM_NO_VERSIONED_SONAME} -DPython3_EXECUTABLE=${Python3_EXECUTABLE} -DPython3_FIND_VIRTUALENV=${Python3_FIND_VIRTUALENV} -DTHEROCK_SOURCE_DIR=${THEROCK_SOURCE_DIR} @@ -70,6 +71,13 @@ set(DEFAULT_CMAKE_ARGS -DCMAKE_INSTALL_LIBDIR=lib -DCMAKE_PROJECT_TOP_LEVEL_INCLUDES=${CMAKE_CURRENT_SOURCE_DIR}/cmake/external_project_include.cmake ) +if(CMAKE_C_VISIBILITY_PRESET) + list(APPEND DEFAULT_CMAKE_ARGS ${CMAKE_C_VISIBILITY_PRESET}) +endif() +if(CMAKE_CXX_VISIBILITY_PRESET) + list(APPEND DEFAULT_CMAKE_ARGS ${CMAKE_CXX_VISIBILITY_PRESET}) +endif() + ################################################################################ # LLVM diff --git a/python_projects/runtime/.gitignore b/python_projects/runtime/.gitignore new file mode 100644 index 0000000..3a701e2 --- /dev/null +++ b/python_projects/runtime/.gitignore @@ -0,0 +1,4 @@ +__pycache__ +build/ +*.egg-info +*.whl diff --git a/python_projects/runtime/README.md b/python_projects/runtime/README.md new file mode 100644 index 0000000..b9310dc --- /dev/null +++ b/python_projects/runtime/README.md @@ -0,0 +1,4 @@ +Welcome to a minimal installation of ROCM runtime libraries. +This package is intended to be taken as a dep from other Python +ecosystem tools which need simple access to AMD-GPU devices without +a full SDK install. diff --git a/python_projects/runtime/pyproject.toml b/python_projects/runtime/pyproject.toml new file mode 100644 index 0000000..2895ac6 --- /dev/null +++ b/python_projects/runtime/pyproject.toml @@ -0,0 +1,11 @@ +[build-system] +requires = [ + "setuptools>=62.3", + "wheel", + "cmake", + # TODO: Re-enable for real. + #"ninja", + # ROCM requires this. + "CppHeaderParser", +] +build-backend = "setuptools.build_meta" diff --git a/python_projects/runtime/setup.py b/python_projects/runtime/setup.py new file mode 100644 index 0000000..8acc3b1 --- /dev/null +++ b/python_projects/runtime/setup.py @@ -0,0 +1,145 @@ +from pathlib import Path +import os +import subprocess +import sys + +from distutils.command.build import build as _build +from setuptools import Extension, find_packages, setup +from setuptools.command.build_ext import build_ext as _build_ext +from setuptools.command.build_py import build_py as _build_py + +try: + from wheel.bdist_wheel import bdist_wheel as _bdist_wheel + + class bdist_wheel(_bdist_wheel): + def finalize_options(self): + # this is a universal, but platform-specific package; a combination + # that wheel does not recognize, thus simply fool it + from distutils.util import get_platform + + self.plat_name = get_platform() + _bdist_wheel.finalize_options(self) + self.root_is_pure = True + +except ImportError: + bdist_wheel = None + + +VERSION = os.getenv("THEROCK_PY_VERSION", "0.1.dev1") +SETUPPY_DIR = Path(__file__).resolve().parent +SOURCE_DIR = SETUPPY_DIR.parent.parent +# Note that setuptools always builds into a "build" directory that +# is a sibling of setup.py, so we just colonize a sub-directory of that +# by default. +CMAKE_BUILD_DIR = os.getenv( + "THEROCK_CMAKE_BUILD_DIR", SETUPPY_DIR / "build" / "cmake-build" +) +CMAKE_INSTALL_DIR = os.getenv( + "THEROCK_CMAKE_INSTALL_DIR", SETUPPY_DIR / "build" / "dist-install" +) + +with open(SETUPPY_DIR / "README.md", "rt") as f: + README = f.read() + + +def getenv_bool(key, default_value="OFF"): + value = os.getenv(key, default_value) + return value.upper() in ["ON", "1", "TRUE"] + + +class CMakeBuildPy(_build_py): + def run(self): + super().run() + print("*****************************", file=sys.stderr) + print("* Building base runtime *", file=sys.stderr) + print("*****************************", file=sys.stderr) + if not getenv_bool("THEROCK_PY_SKIP_CMAKE"): + self.cmake_build() + else: + print("Skipping CMake build because THEROCK_PY_SKIP_CMAKE is set") + self.cmake_install() + + def cmake_build(self): + subprocess.check_call(["cmake", "--version"]) + install_dir = CMAKE_INSTALL_DIR + cmake_args = [ + f"-S{SOURCE_DIR}", + f"-B{CMAKE_BUILD_DIR}", + "-GNinja", + "--log-level=VERBOSE", + f"-DCMAKE_BUILD_TYPE=Release", + f"-DCMAKE_INSTALL_PREFIX={str(install_dir)}", + # Versioned SONAMEs get duplicated in Python wheels. And this use + # case is really aimed at dynamically loading of a hermetic runtime + # library. + "-DCMAKE_PLATFORM_NO_VERSIONED_SONAME=ON", + # TODO: Also enable visibility=hidden preset and other distribution + # armor. + ] + subprocess.check_call(["cmake"] + cmake_args, cwd=SOURCE_DIR) + subprocess.check_call( + ["cmake", "--build", CMAKE_BUILD_DIR], cwd=CMAKE_BUILD_DIR + ) + + def cmake_install(self): + (CMAKE_INSTALL_DIR / "__init__.py").touch() + with open(CMAKE_INSTALL_DIR / "version.py", "wt") as f: + f.write(f"py_version = '{VERSION}'\n") + # TODO: Add ROCM version, etc + for component in ["amdgpu-runtime"]: + subprocess.check_call( + ["cmake", "--install", CMAKE_BUILD_DIR, "--component", component], + cwd=CMAKE_BUILD_DIR, + ) + + +class CustomBuild(_build): + def run(self): + self.run_command("build_py") + self.run_command("build_ext") + self.run_command("build_scripts") + + +packages = find_packages(where=".") +print("Found packages:", packages) + +CMAKE_INSTALL_DIR.mkdir(parents=True, exist_ok=True) +CMAKE_BUILD_DIR.mkdir(parents=True, exist_ok=True) + +setup( + name=f"TheRock-runtime", + version=f"{VERSION}", # TODO: Get from env var. + author="TheRock Authors", + author_email="stdin@nod.ai", + description="Minimal ROCM runtime components", + long_description=README, + long_description_content_type="text/markdown", + license="Apache-2.0", + classifiers=[ + "Development Status :: 3 - Alpha", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + ], + url="https://github.com/nod-ai/TheRock", + python_requires=">=3.6", + cmdclass={ + "build": CustomBuild, + "build_py": CMakeBuildPy, + "bdist_wheel": bdist_wheel, + }, + zip_safe=False, + packages=["therock", "_therock"], + package_dir={ + "": ".", + "_therock": "build/dist-install", + }, + # Matching the native extension as a data file keeps setuptools from + # "building" it (i.e. turning it into a static binary). + package_data={ + "_therock": [ + "**/*", + ], + }, + entry_points={}, + install_requires=[], +) diff --git a/python_projects/runtime/therock/__init__.py b/python_projects/runtime/therock/__init__.py new file mode 100644 index 0000000..395382d --- /dev/null +++ b/python_projects/runtime/therock/__init__.py @@ -0,0 +1,37 @@ +from pathlib import Path +import platform + +__all__ = [ + "get_dist_dir", + "get_library_dir", + "get_hip_runtime_library", +] + + +def _get_library_parts() -> tuple[str, str]: + s = platform.system() + if s == "Linux": + return "lib", ".so" + elif s == "Windows": + return "", ".dll" + elif s == "Darwin": + return "lib", ".dylib" + + +lib_prefix, lib_suffix = _get_library_parts() + + +def get_dist_dir() -> str: + import _therock + + location = _therock.__file__ + assert location, "Could not find physical location for therock runtime" + return str(Path(location).resolve().parent) + + +def get_library_dir() -> str: + return str(Path(get_dist_dir()) / "lib") + + +def get_hip_runtime_library() -> str: + return str(Path(get_library_dir()) / f"{lib_prefix}amdhip64{lib_suffix}") diff --git a/python_projects/runtime/therock/test_dist.py b/python_projects/runtime/therock/test_dist.py new file mode 100644 index 0000000..26a6220 --- /dev/null +++ b/python_projects/runtime/therock/test_dist.py @@ -0,0 +1,53 @@ +import ctypes +from pathlib import Path +import unittest + +import therock + + +class TestVersion(unittest.TestCase): + def testVersion(self): + import _therock.version as v + + py_version = v.py_version + print("Found py_version =", py_version) + self.assertTrue(py_version) + + +class TestLocations(unittest.TestCase): + def testLibDir(self): + lib_dir = therock.get_library_dir() + print(lib_dir) + self.assertTrue(Path(lib_dir).is_dir(), msg=lib_dir) + + def testHipRuntimeLibrary(self): + rtl = Path(therock.get_hip_runtime_library()) + self.assertTrue(rtl.is_file, msg=rtl) + + +class TestLoadRuntimeLibrary(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.dylib = ctypes.CDLL(therock.get_hip_runtime_library()) + cls.dylib.hipRuntimeGetVersion.restype = ctypes.c_int + cls.dylib.hipRuntimeGetVersion.argtypes = [ctypes.POINTER(ctypes.c_int)] + cls.dylib.hipDriverGetVersion.restype = ctypes.c_int + cls.dylib.hipDriverGetVersion.argtypes = [ctypes.POINTER(ctypes.c_int)] + + def testRuntimeVersion(self): + v = ctypes.c_int() + rc = self.dylib.hipRuntimeGetVersion(v) + print("RuntimeVersion =", v.value) + self.assertEqual(rc, 0) + self.assertGreaterEqual(v.value, 60000000) + + def testDriverVersion(self): + v = ctypes.c_int() + rc = self.dylib.hipDriverGetVersion(v) + print("DriverVersion =", v.value) + self.assertEqual(rc, 0) + self.assertGreaterEqual(v.value, 60000000) + + +if __name__ == "__main__": + unittest.main()