-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
3845d11
commit 4ee722d
Showing
7 changed files
with
262 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,4 @@ | ||
__pycache__ | ||
build/ | ||
*.egg-info | ||
*.whl |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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. |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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="[email protected]", | ||
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=[], | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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}") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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() |