nixos-lib/hardware/amd/rocm.nix

{
  config,
  pkgs,
  lib,
  ...
}:
with lib;
let
  cfg = config.aviallon.hardware.amd;
  localCfg = config.aviallon.hardware.amd.rocm;
  devCfg = config.aviallon.developer;
  generalCfg = config.aviallon.general;

  gfxToCompatibleMap = {
    gfx900 = "9.0.0";
    gfx902 = "9.0.0";
    gfx904 = "9.0.0";
    gfx909 = "9.0.0";
    gfx90c = "9.0.0";
    gfx1011 = "10.1.0";
    gfx1012 = "10.1.0";
    gfx1031 = "10.3.0";
    gfx1032 = "10.3.0";
    gfx1033 = "10.3.0";
    gxf1034 = "10.3.0";
    gxf1035 = "10.3.0";
    gxf1036 = "10.3.0";
  };

  /*
    autoDetectGPU = pkgs: pkgs.callPackage (
      { runCommandLocal,
        gnugrep,
        rocmPackages,
      }: runCommandLocal "hsa-version" { nativeBuildInputs = [ gnugrep rocmPackages.rocminfo ]; } ''
          set +e
          mkdir -p $out/
          echo "Computing HSA version" &>/dev/stderr
          ls -l /dev/kfd
          rocminfo &>/dev/stderr
          rocminfo | grep --only-matching --perl-regexp '^\s*Name:\s+\Kgfx[0-9a-f]+' | tee $out/output
         ''
    ) { };
  */

  gfxToCompatible =
    gfxISA: if (hasAttr gfxISA gfxToCompatibleMap) then (getAttr gfxISA gfxToCompatibleMap) else "";
in
{

  options.aviallon.hardware.amd.rocm = {
    enable = (mkEnableOption "ROCm configuration") // {
      default = true;
    };
    gfxISA = mkOption {
      description = "What is the GFX ISA of your system. Leave blank if you have several GPUs of incompatible ISAs";
      default = "";
      example = "gfx902";
      type = types.str;
    };
    gpuTargets = mkOption {
      description = "Override supported GPU ISAs in some ROCm packages.";
      default = [
        "803"
        "900"
        "906:xnack-"
        "908:xnack-"
        "90a:xnack+"
        "90a:xnack-"
        "940"
        "941"
        "942"
        "1010"
        "1012"
        "1030"
        "1031"
        "1100"
        "1101"
        "1102"
      ];
      example = [
        "900"
        "1031"
      ];
      type = with types; nullOr (listOf str);
    };
  };

  config = mkIf (cfg.enable && localCfg.enable) {
    environment.systemPackages =
      with pkgs;
      [
        rocmPackages.rocm-smi
        #rocmPackages.meta.rocm-ml-libraries
        #rocmPackages.meta.rocm-hip-runtime

        #pkgs.autoDetectGPU
      ]
      ++ optionals devCfg.enable [
        rocmPackages.rocminfo
      ];

    #systemd.tmpfiles.rules = [
    #  "L+    /opt/rocm/hip   -    -    -     -    ${pkgs.rocmPackages.meta.rocm-hip-runtime}"
    #"L+    /tmp/hsa-version - - - - ${pkgs.autoDetectGPU}"
    #];

    environment.variables = {
      ROC_ENABLE_PRE_VEGA = "1"; # Enable OpenCL with Polaris GPUs
    }
    // (mkIf (gfxToCompatible cfg.rocm.gfxISA != "") {
      HSA_OVERRIDE_GFX_VERSION = gfxToCompatible cfg.rocm.gfxISA;
    });

    # Make rocblas and rocfft work
    nix.settings.extra-sandbox-paths = [
      "/dev/kfd?"
      "/sys/devices/virtual/kfd?"
      "/dev/dri/renderD128?"
    ];

    nix.settings.substituters = [ "https://nixos-rocm.cachix.org" ];
    nix.settings.trusted-public-keys = [
      "nixos-rocm.cachix.org-1:VEpsf7pRIijjd8csKjFNBGzkBqOmw8H9PRmgAq14LnE="
    ];

    nixpkgs.config.rocmSupport = true;

    nixpkgs.overlays = mkIf (!isNull localCfg.gpuTargets) (mkBefore [
      (final: prev: {
        #rocmPackages_5 = final.rocmPackages;
        rocmPackages = prev.rocmPackages // {
          clr = prev.rocmPackages.clr.override {
            localGpuTargets = lib.forEach localCfg.gpuTargets (target: "gfx${target}");
          };
          rocdbgapi = prev.rocmPackages.rocdbgapi.override { buildDocs = false; };
          # (oldAttrs: {
          #  passthru = oldAttrs.passthru // {
          #    # We cannot use this for each ROCm library, as each defines their own supported targets
          #    # See: https://github.com/ROCm/ROCm/blob/77cbac4abab13046ee93d8b5bf410684caf91145/README.md#library-target-matrix
          #    gpuTargets = lib.forEach localCfg.gpuTargets (target: "gfx${target}");
          #  };
          #});

          #rocblas = prev.rocmPackages.rocblas.override {
          #  gpuTargets = lib.forEach localCfg.gpuTargets (target: "gfx${target}");
          #};
        };
      })
    ]);
  };
}