Numbus/modules/server/hardware/pcie-coral.nix

{ config, lib, pkgs, ... }:

let
  cfg = config.numbus-server.hardware.pcie-coral;

  gasket-driver = { stdenv, lib, fetchFromGitHub, kernel }: stdenv.mkDerivation rec {
    pname = "gasket";
    version = "1.0-18";

    src = fetchFromGitHub {
      owner = "google";
      repo = "gasket-driver";
      rev = "97aeba584efd18983850c36dcf7384b0185284b3";
      sha256 = "pJwrrI7jVKFts4+bl2xmPIAD01VKFta2SRuElerQnTo=";
    };

    makeFlags = [
      "-C"
      "${kernel.dev}/lib/modules/${kernel.modDirVersion}/build"
      "M=$(PWD)"
    ];
    buildFlags = [ "modules" ];

    installFlags = [ "INSTALL_MOD_PATH=${placeholder "out"}" ];
    installTargets = [ "modules_install" ];

    sourceRoot = "source/src";
    hardeningDisable = [ "pic" "format" ];
    nativeBuildInputs = kernel.moduleBuildDependencies;

    meta = with lib; {
      description = "The Coral Gasket Driver allows usage of the Coral EdgeTPU on Linux systems.";
      homepage = "https://github.com/google/gasket-driver";
      license = licenses.gpl2;
      maintainers = [ maintainers.kylehendricks ];
      platforms = platforms.linux;
    };
  };

  libedgetpu-pkg = { stdenv, lib, fetchFromGitHub, libusb1, abseil-cpp, flatbuffers, xxd }:
    let
      flatbuffers_1_12 = flatbuffers.overrideAttrs (oldAttrs: rec {
        version = "1.12.0";
        NIX_CFLAGS_COMPILE = "-Wno-error=class-memaccess -Wno-error=maybe-uninitialized";
        cmakeFlags = (oldAttrs.cmakeFlags or []) ++ ["-DFLATBUFFERS_BUILD_SHAREDLIB=ON"];
        NIX_CXXSTDLIB_COMPILE = "-std=c++17";
        configureFlags = (oldAttrs.configureFlags or []) ++ ["--enable-shared"];
        src = fetchFromGitHub {
          owner = "google";
          repo = "flatbuffers";
          rev = "v${version}";
          sha256 = "sha256-L1B5Y/c897Jg9fGwT2J3+vaXsZ+lfXnskp8Gto1p/Tg=";
        };
      });

    in stdenv.mkDerivation rec {
      pname = "libedgetpu";
      version = "grouper";

      src = fetchFromGitHub {
        owner = "google-coral";
        repo = pname;
        rev = "release-${version}";
        sha256 = "sha256-73hwItimf88Iqnb40lk4ul/PzmCNIfdt6Afi+xjNiBE=";
      };

      makeFlags = ["-f" "makefile_build/Makefile" "libedgetpu" ];

      buildInputs = [
        libusb1
        abseil-cpp
        flatbuffers_1_12
      ];

      nativeBuildInputs = [
        xxd
      ];

      NIX_CXXSTDLIB_COMPILE = "-std=c++17";

      TFROOT = "${fetchFromGitHub {
        owner = "tensorflow";
        repo = "tensorflow";
        rev = "v2.7.4";
        sha256 = "sha256-liDbUAdaVllB0b74aBeqNxkYNu/zPy7k3CevzRF5dk0=";
      }}";

      enableParallelBuilding = false;

      installPhase = ''
        mkdir -p $out/lib
        cp out/direct/k8/libedgetpu.so.1.0 $out/lib
        ln -s $out/lib/libedgetpu.so.1.0 $out/lib/libedgetpu.so.1
        mkdir -p $out/lib/udev/rules.d
        cp debian/edgetpu-accelerator.rules $out/lib/udev/rules.d/99-edgetpu-accelerator.rules
      '';
    };

  gasket = config.boot.kernelPackages.callPackage gasket-driver {};
  libedgetpu = pkgs.callPackage libedgetpu-pkg {};
in

{
  options.numbus-server.hardware.pcie-coral = lib.mkEnableOption "PCIe Coral TPU support";

  config = lib.mkIf cfg.enable {
    services.udev.packages = [ libedgetpu ];
    users.groups.plugdev = {};
    boot.extraModulePackages = [ gasket ];
  };
}