Flatten build variants to build/<variant>

lib/torch-extension/arch.nix

@@ -13,6 +13,7 @@
   cuda_nvcc,
   get-kernel-check,
   kernel-abi-check,
+  kernel-layout-check,
   ninja,
   python3,
   remove-bytecode-hook,
@@ -79,7 +80,12 @@ in
 stdenv.mkDerivation (prevAttrs: {
   name = "${extensionName}-torch-ext";
 
-  inherit doAbiCheck nvccThreads src;
+  inherit
+    doAbiCheck
+    extensionName
+    nvccThreads
+    src
+    ;
 
   # Generate build files.
   postPatch = ''
@@ -123,10 +129,11 @@ stdenv.mkDerivation (prevAttrs: {
   '';
 
   nativeBuildInputs = [
-    kernel-abi-check
     cmake
     ninja
     build2cmake
+    kernel-abi-check
+    kernel-layout-check
     remove-bytecode-hook
   ]
   ++ lib.optionals doGetKernelCheck [
@@ -223,28 +230,31 @@ 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" {} \;
   '';
 
   doInstallCheck = true;
 
-  getKernelCheck = extensionName;
-
   # We need access to the host system on Darwin for the Metal compiler.
   __noChroot = metalSupport;
 

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

@@ -4,6 +4,7 @@
 
   build2cmake,
   get-kernel-check,
+  kernel-layout-check,
   remove-bytecode-hook,
   torch,
 }:
@@ -23,14 +24,15 @@
 stdenv.mkDerivation (prevAttrs: {
   name = "${extensionName}-torch-ext";
 
-  inherit src;
+  inherit extensionName src;
 
   # Add Torch as a dependency, so that devshells for universal kernels
   # also get torch as a build input.
   buildInputs = [ torch ];
 
   nativeBuildInputs = [
     build2cmake
+    kernel-layout-check
     remove-bytecode-hook
   ]
   ++ lib.optionals doGetKernelCheck [
@@ -48,10 +50,10 @@ 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;
-
-  getKernelCheck = extensionName;
 })

overlay.nix

@@ -9,6 +9,8 @@ final: prev: {
 
   kernel-abi-check = prev.callPackage ./pkgs/kernel-abi-check { };
 
+  kernel-layout-check = prev.callPackage ./pkgs/kernel-layout-check { };
+
   rewrite-nix-paths-macho = prev.callPackage ./pkgs/rewrite-nix-paths-macho { };
 
   remove-bytecode-hook = prev.callPackage ./pkgs/remove-bytecode-hook { };

pkgs/get-kernel-check/get-kernel-check-hook.sh

@@ -3,35 +3,39 @@
 echo "Sourcing get-kernel-check-hook.sh"
 
 _getKernelCheckHook() {
-  if [ ! -z "${getKernelCheck}" ]; then
-    echo "Checking loading kernel with get_kernel"
-    echo "Check whether the kernel can be loaded with get-kernel: ${getKernelCheck}"
-
-    # We strip the full library paths from the extension. Unfortunately,
-    # in a Nix environment, the library dependencies cannot be found
-    # anymore. So we have to add the Torch library directory to the
-    # dynamic linker path to get it to pick it up.
-    if [ $(uname -s) == "Darwin" ]; then
-      TORCH_DIR=$(python -c "from pathlib import Path; import torch; print(Path(torch.__file__).parent)")
-      export DYLD_LIBRARY_PATH="${TORCH_DIR}/lib:${DYLD_LIBRARY_PATH}"
-    fi
-
-    TMPDIR=$(mktemp -d -t test.XXXXXX) || exit 1
-    trap "rm -rf '$TMPDIR'" EXIT
-
-    # Some kernels want to write stuff (especially when they use Triton).
-    HOME=$(mktemp -d -t test.XXXXXX) || exit 1
-    trap "rm -rf '$HOME'" EXIT
-
-    # Emulate the bundle layout that kernels expects. This even works
-    # for universal kernels, since kernels checks the non-universal
-    # path first.
-    BUILD_VARIANT=$(python -c "from kernels.utils import build_variant; print(build_variant())")
-    mkdir -p "${TMPDIR}/build"
-    ln -s "$out" "${TMPDIR}/build/${BUILD_VARIANT}"
-
-    python -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('${TMPDIR}'), '${getKernelCheck}')"
+  echo "Checking loading kernel with get_kernel"
+
+  if [ -z ${extensionName+x} ]; then
+    echo "extensionName must be set in derivation"
+    exit 1
+  fi
+
+  echo "Check whether the kernel can be loaded with get-kernel: ${extensionName}"
+
+  # We strip the full library paths from the extension. Unfortunately,
+  # in a Nix environment, the library dependencies cannot be found
+  # anymore. So we have to add the Torch library directory to the
+  # dynamic linker path to get it to pick it up.
+  if [ $(uname -s) == "Darwin" ]; then
+    TORCH_DIR=$(python -c "from pathlib import Path; import torch; print(Path(torch.__file__).parent)")
+    export DYLD_LIBRARY_PATH="${TORCH_DIR}/lib:${DYLD_LIBRARY_PATH}"
   fi
+
+  TMPDIR=$(mktemp -d -t test.XXXXXX) || exit 1
+  trap "rm -rf '$TMPDIR'" EXIT
+
+  # Some kernels want to write stuff (especially when they use Triton).
+  HOME=$(mktemp -d -t test.XXXXXX) || exit 1
+  trap "rm -rf '$HOME'" EXIT
+
+  # Emulate the bundle layout that kernels expects. This even works
+  # for universal kernels, since kernels checks the non-universal
+  # path first.
+  BUILD_VARIANT=$(python -c "from kernels.utils import build_variant; print(build_variant())")
+  mkdir -p "${TMPDIR}/build"
+  ln -s "$out" "${TMPDIR}/build/${BUILD_VARIANT}"
+
+  python -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('${TMPDIR}'), '${extensionName}')"
 }
 
 postInstallCheckHooks+=(_getKernelCheckHook)

pkgs/kernel-abi-check/kernel-abi-check-hook.sh

@@ -5,7 +5,7 @@ _checkAbiHook() {
     echo "Skipping ABI check"
   else
     echo "Checking of ABI compatibility"
-    find "$out/${extensionName}" -name '*.so' -print0 | \
+    find "$out/" -name '*.so' -print0 | \
       xargs -0 -n1 kernel-abi-check
   fi
 }

pkgs/kernel-layout-check/default.nix

@@ -0,0 +1,5 @@
+{ makeSetupHook, python3 }:
+
+makeSetupHook {
+  name = "kernel-layout-check-hook";
+} ./kernel-layout-check-hook.sh

pkgs/kernel-layout-check/kernel-layout-check-hook.sh

@@ -0,0 +1,27 @@
+#!/bin/sh
+
+echo "Sourcing kernel-layout-check-hook.sh"
+
+kernelLayoutCheckHook() {
+  echo "Checking kernel layout"
+
+  if [ -z ${extensionName+x} ]; then
+    echo "extensionName must be set in derivation"
+    exit 1
+  fi
+
+  if [ ! -f source/torch-ext/${extensionName}/__init__.py ]; then
+    echo "Python module at source/torch-ext/${extensionName} must contain __init__.py"
+    exit 1
+  fi
+
+  # TODO: remove once the old location is removed from kernels.
+  if [ -e source/torch-ext/${extensionName}/${extensionName} ]; then
+    echo "Python module at source/torch-ext/${extensionName} must not have ${extensionName} file or directory."
+    exit 1
+  fi
+}
+
+if [ -z "${dontCheckLayout-}" ]; then
+  postUnpackHooks+=(kernelLayoutCheckHook)
+fi