Add some kernel layout checks

- Validate that the kernel does not have a file/directory
  `torch-ext/<mod_name>/<mod_name>` since we need that directory for the
  compat module.
- Validate that the kernel has `torch-ext/<mod_name>/__init__.py`. This
  is also verified by get-kernel-check, but that only fires *after* the
  kernel is already compiled.

Author: danieldk

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 [
@@ -248,8 +255,6 @@ stdenv.mkDerivation (prevAttrs: {
 
   doInstallCheck = true;
 
-  getKernelCheck = extensionName;
-
   # We need access to the host system on Darwin for the Metal compiler.
   __noChroot = metalSupport;
 

lib/torch-extension/no-arch.nix

@@ -4,6 +4,7 @@
 
   build2cmake,
   get-kernel-check,
+  kernel-layout-check,
   torch,
 }:
 
@@ -22,14 +23,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
   ]
   ++ lib.optionals doGetKernelCheck [
     get-kernel-check
@@ -52,6 +54,4 @@ stdenv.mkDerivation (prevAttrs: {
   '';
 
   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-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