From 1e31b1599c8434405f018fae04c0bec679cabdfe Mon Sep 17 00:00:00 2001
From: Max <max@privatevoid.net>
Date: Wed, 12 Apr 2023 18:23:12 +0200
Subject: [PATCH] koboldai: apply ROCm workarounds

---
 projects/koboldai/package.nix | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/projects/koboldai/package.nix b/projects/koboldai/package.nix
index 5d1ed5e..c40b96d 100644
--- a/projects/koboldai/package.nix
+++ b/projects/koboldai/package.nix
@@ -7,6 +7,7 @@
 , runCommand
 , tmpDir ? "/tmp/nix-koboldai"
 , stateDir ? "$HOME/.koboldai/state"
+, libdrm
 }:
 let
   overrides = {
@@ -87,6 +88,17 @@ let
     lupa
     memcached
   ]);
+
+  # See note about consumer GPUs:
+  # https://docs.amd.com/bundle/ROCm-Deep-Learning-Guide-v5.4.3/page/Troubleshooting.html
+  rocmInit = ''
+    if [ ! -e /tmp/nix-pytorch-rocm___/amdgpu.ids ]
+    then
+        mkdir -p /tmp/nix-pytorch-rocm___
+        ln -s ${libdrm}/share/libdrm/amdgpu.ids /tmp/nix-pytorch-rocm___/amdgpu.ids
+    fi
+    export HSA_OVERRIDE_GFX_VERSION=''${HSA_OVERRIDE_GFX_VERSION-'10.3.0'}
+  '';
 in
 (writeShellScriptBin "koboldai" ''
   if [ -d "/usr/lib/wsl/lib" ]
@@ -102,6 +114,7 @@ in
   ln -s ${stateDir}/models/   ${tmpDir}/models
   ln -s ${stateDir}/settings/ ${tmpDir}/settings
   ln -s ${stateDir}/userscripts/ ${tmpDir}/userscripts
+  ${lib.optionalString (aipython3.torch.rocmSupport or false) rocmInit}
   ${koboldPython}/bin/python ${patchedSrc}/aiserver.py $@
 '').overrideAttrs
   (_: {