Flatten build variants to build/<variant>

Prior to this change, kernels were stored in `build/<variant>/<extname>`.
However, this was fragile because the extension name had to correspond to
the repository name. This change flattens kernels to be stored inside
`build/<variant>`.

For compatibility with older versions of kernels, we add a module
`build/<variant>/<extname>` that loads `build/<variant>`, this
compatibility module will removed when the `kernels` update has been
around for a while.

Author: danieldk

lib/torch-extension/arch.nix

@@ -223,21 +223,26 @@ stdenv.mkDerivation (prevAttrs: {
   postInstall = ''
     (
       cd ..
-      cp -r torch-ext/${extensionName} $out/
+      cp -r torch-ext/${extensionName}/* $out/
     )
-    cp $out/_${extensionName}_*/* $out/${extensionName}
-    rm -rf $out/_${extensionName}_*
+    mv $out/_${extensionName}_*/* $out/
+    rm -d $out/_${extensionName}_${rev}
+
+    # Set up a compatibility module for older kernels versions, remove when
+    # the updated kernels has been around for a while.
+    mkdir $out/${extensionName}
+    cp ${./compat.py} $out/${extensionName}/__init__.py
   ''
   + (lib.optionalString (stripRPath && stdenv.hostPlatform.isLinux)) ''
-    find $out/${extensionName} -name '*.so' \
+    find $out/ -name '*.so' \
       -exec patchelf --set-rpath "" {} \;
   ''
   + (lib.optionalString (stripRPath && stdenv.hostPlatform.isDarwin)) ''
-    find $out/${extensionName} -name '*.so' \
+    find $out/ -name '*.so' \
       -exec rewrite-nix-paths-macho {} \;
 
     # Stub some rpath.
-    find $out/${extensionName} -name '*.so' \
+    find $out/ -name '*.so' \
       -exec install_name_tool -add_rpath "@loader_path/lib" {} \;
   '';
 

lib/torch-extension/compat.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

lib/torch-extension/no-arch.nix

@@ -46,7 +46,9 @@ stdenv.mkDerivation (prevAttrs: {
 
   installPhase = ''
     mkdir -p $out
-    cp -r torch-ext/${extensionName} $out/
+    cp -r torch-ext/${extensionName}/* $out/
+    mkdir $out/${extensionName}
+    cp ${./compat.py} $out/${extensionName}/__init__.py
   '';
 
   doInstallCheck = true;