From 0682a7717108855f7d92b207f0ce8e15f3d230ec Mon Sep 17 00:00:00 2001
From: RisingPrism <git@risingprism.com>
Date: Sat, 3 Oct 2020 20:25:26 -0400
Subject: [PATCH] Make scripts work on more dms and more gpus

---
 polaris-vfio-startup.sh  | 90 ++++++++++++++++++++++++++++++++++++----
 polaris-vfio-teardown.sh | 32 ++++++++++++--
 2 files changed, 109 insertions(+), 13 deletions(-)

diff --git a/polaris-vfio-startup.sh b/polaris-vfio-startup.sh
index d50ed5a..d39560c 100755
--- a/polaris-vfio-startup.sh
+++ b/polaris-vfio-startup.sh
@@ -2,19 +2,89 @@
 # Helpful to read output when debugging
 set -x
 
-# Stop display manager (KDE specific)
-systemctl stop sddm.service
+long_delay=10
+medium_delay=5
+short_delay=1
+echo "Beginning of startup!"
 
-# Unbind VTconsoles
-echo 0 > /sys/class/vtconsole/vtcon0/bind
+# Stop currently running display manager
+if test -e "/tmp/vfio-store-display-manager" ; then
+    rm -f /tmp/vfio-store-display-manager
+fi
+if systemctl is-active --quiet sddm.service ; then
+    echo sddm.service >> /tmp/vfio-store-display-manager
+    systemctl stop sddm.service
+fi
+while systemctl is-active --quiet sddm.service ; do
+    sleep "${short_delay}"
+done
+if systemctl is-active --quiet gdm.service ; then
+    echo gdm.service >> /tmp/vfio-store-display-manager
+    systemctl stop gdm.service
+fi
+if systemctl is-active --quiet lightdm.service ; then
+    echo lightdm.service >> /tmp/vfio-store-display-manager
+    systemctl stop lightdm.service
+fi
+if systemctl is-active --quiet lxdm.service ; then
+    echo lxdm.service >> /tmp/vfio-store-display-manager
+    systemctl stop lxdm.service
+fi
+if systemctl is-active --quiet xdm.service ; then
+    echo xdm.service >> /tmp/vfio-store-display-manager
+    systemctl stop xdm.service
+fi
 
-# Unbind EFI-Framebuffer
-echo efi-framebuffer.0 > /sys/bus/platform/drivers/efi-framebuffer/unbind
+# Unbind VTconsoles if currently bound
+if test -e "/sys/class/vtconsole/vtcon0/bind" ; then
+    echo 0 > /sys/class/vtconsole/vtcon0/bind
+    sleep "${long_delay}"
+fi
+if test -e "/sys/class/vtconsole/vtcon1/bind" ; then
+    echo 0 > /sys/class/vtconsole/vtcon1/bind
+    sleep "${long_delay}"
+fi
 
-sleep 5
+#Unbind EFI-Framebuffer if currently bound
+if test -e "/sys/bus/platform/drivers/efi-framebuffer/unbind" ; then
+    echo efi-framebuffer.0 > /sys/bus/platform/drivers/efi-framebuffer/unbind
+    sleep "${medium_delay}"
+else
+    echo "Could not find framebuffer to unload!"
+fi
+
+# Unload loaded GPU drivers
+if test -e "/tmp/vfio-loaded-gpu-modules" ; then
+    rm -f /tmp/vfio-loaded-gpu-modules
+fi
+if lsmod | grep amdgpu &> /dev/null ; then
+    modprobe -r amdgpu
+    echo amdgpu >> /tmp/vfio-loaded-gpu-modules
+fi
+while lsmod | grep amdgpu &> /dev/null ; do
+    sleep 1
+done
+if lsmod | grep nvidia_drm &> /dev/null ; then
+    modprobe -r nvidia_drm
+    echo nvidia_drm >> /tmp/vfio-loaded-gpu-modules
+fi
+if lsmod | grep nvidia_modeset &> /dev/null ; then
+    modprobe -r nvidia_modeset
+    echo nvidia_modeset >> /tmp/vfio-loaded-gpu-modules
+fi
+if lsmod | grep nvidia_uvm &> /dev/null ; then
+    modprobe -r nvidia_uvm
+    echo nvidia_uvm >> /tmp/vfio-loaded-gpu-modules
+fi
+if lsmod | grep nvidia &> /dev/null ; then
+    modprobe -r nvidia
+    echo nvidia >> /tmp/vfio-loaded-gpu-modules
+fi
+if lsmod | grep ipmi_devintf &> /dev/null ; then
+    modprobe -r ipmi_devintf
+    echo ipmi_devintf >> /tmp/vfio-loaded-gpu-modules
+fi
 
-# Unload AMD drivers
-modprobe -r amdgpu
 
 # Unbind the GPU from display driver
 virsh nodedev-detach pci_0000_01_00_0
@@ -22,3 +92,5 @@ virsh nodedev-detach pci_0000_01_00_1
 
 # Load VFIO kernel module
 modprobe vfio-pci
+
+echo "End of startup!"
\ No newline at end of file
diff --git a/polaris-vfio-teardown.sh b/polaris-vfio-teardown.sh
index b94d887..29dde10 100755
--- a/polaris-vfio-teardown.sh
+++ b/polaris-vfio-teardown.sh
@@ -1,6 +1,8 @@
 #!/bin/bash
 set -x
 
+echo "Beginning of teardown!"
+
 # Unload VFIO-PCI Kernel Driver
 modprobe -r vfio-pci
 modprobe -r vfio_iommu_type1
@@ -11,13 +13,35 @@ virsh nodedev-reattach pci_0000_01_00_1
 virsh nodedev-reattach pci_0000_01_00_0
 
 # Rebind VT consoles
-echo 1 > /sys/class/vtconsole/vtcon0/bind
+if test -e "/sys/class/vtconsole/vtcon0/bind" ; then
+    echo 1 > /sys/class/vtconsole/vtcon0/bind
+fi
+if test -e "/sys/class/vtconsole/vtcon1/bind" ; then
+    echo 1 > /sys/class/vtconsole/vtcon1/bind
+fi
+
+# Hack that magically makes nvidia gpus work :)
+if command -v nvidia-xconfig ; then
+    nvidia-xconfig --query-gpu-info > /dev/null 2>&1
+fi
 
 # Re-Bind EFI-Framebuffer
-echo "efi-framebuffer.0" > /sys/bus/platform/drivers/efi-framebuffer/bind
+if test -e "/sys/bus/platform/drivers/efi-framebuffer/bind" ; then
+    echo "efi-framebuffer.0" > /sys/bus/platform/drivers/efi-framebuffer/bind
+else
+    echo "Could not find framebuffer to bind!"
+fi
 
 #Load amd driver
-modprobe amdgpu
+input="/tmp/vfio-loaded-gpu-modules"
+while read gpuModule; do
+  modprobe "$gpuModule"
+done < "$input"
 
 # Restart Display Manager
-systemctl start sddm.service
+input="/tmp/vfio-store-display-manager"
+while read displayManager; do
+  systemctl start "$displayManager"
+done < "$input"
+
+echo "End of teardown!"