refactor: move compiled kernel .so files into extension/lib

- CUDAExtension module names changed from astrai.extension.<name> to astrai.extension.lib.<name>
- Compiled .so files now land in astrai/extension/lib/ instead of alongside Python source
- loader.py imports from .lib.<name> subpackage
- Add astrai/extension/lib/__init__.py to make lib a proper package
- Separates compiled artifacts from Python source for cleaner directory structure
This commit is contained in:
2026-07-31 15:36:32 +08:00
parent 3e67b4f88d
commit 9f83d982ec
3 changed files with 4 additions and 3 deletions
+1
View File
@@ -0,0 +1 @@
"""Compiled CUDA kernel modules (``*.so``) live here, kept separate from Python source."""
+1 -1
View File
@@ -18,7 +18,7 @@ _modules: dict[str, object] = {}
for _name in KERNEL_NAMES: for _name in KERNEL_NAMES:
try: try:
_mod = importlib.import_module(f".{_name}", package=__package__) _mod = importlib.import_module(f".lib.{_name}", package=__package__)
_available[_name] = True _available[_name] = True
_modules[_name] = _mod _modules[_name] = _mod
except ImportError: except ImportError:
+2 -2
View File
@@ -7,7 +7,7 @@ from setuptools import setup
from setuptools.command.build_ext import build_ext as _build_ext from setuptools.command.build_ext import build_ext as _build_ext
sys.path.insert(0, str(Path(__file__).parent)) sys.path.insert(0, str(Path(__file__).parent))
os.makedirs("astrai/extension", exist_ok=True) os.makedirs("astrai/extension/lib", exist_ok=True)
def _should_build(): def _should_build():
@@ -57,7 +57,7 @@ if _should_build():
for name, info in REGISTRY.items(): for name, info in REGISTRY.items():
ext_modules.append( ext_modules.append(
CUDAExtension( CUDAExtension(
f"astrai.extension.{name}", f"astrai.extension.lib.{name}",
info["sources"], info["sources"],
extra_compile_args={ extra_compile_args={
"cxx": info["cxx_flags"], "cxx": info["cxx_flags"],