From 349dfbfbd72c53e0d7f5aa13175f6608a0a1749a Mon Sep 17 00:00:00 2001
From: Joji Mekkattuparamban <jojim@nvidia.com>
Date: Thu, 27 Jun 2024 16:11:24 -0700
Subject: [PATCH] feat: vfio-manager graphics mode

Signed-off-by: Joji Mekkattuparamban <jojim@nvidia.com>
---
 scripts/vfio-manage | 54 ++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 49 insertions(+), 5 deletions(-)

diff --git a/scripts/vfio-manage b/scripts/vfio-manage
index 4c55672a..de6994c7 100755
--- a/scripts/vfio-manage
+++ b/scripts/vfio-manage
@@ -85,6 +85,12 @@ unbind_device() {
 
     echo "unbinding device $gpu"
     unbind_from_driver $gpu
+    #for graphics mode, we need to unbind the auxiliary device as well
+    aux_dev=$(get_graphics_aux_dev "$gpu")
+    if [ "$aux_dev" != "NONE" ]; then
+      echo "gpu $gpu is in graphics mode aux_dev $aux_dev"
+      unbind_from_driver "$aux_dev"
+    fi
 }
 
 unbind_all() {
@@ -97,13 +103,9 @@ unbind_all() {
     done
 }
 
-bind_device() {
+bind_pci_device() {
     local gpu=$1
 
-    if ! is_nvidia_gpu_device $gpu; then
-        return 0
-    fi
-
     if ! is_bound_to_vfio $gpu; then
       unbind_from_other_driver $gpu
       echo "binding device $gpu"
@@ -114,6 +116,48 @@ bind_device() {
     fi
 }
 
+get_graphics_aux_dev() {
+    local gpu=$1
+    device_class_file=$(readlink -f "/sys/bus/pci/devices/$gpu/class")
+    device_class=$(cat "$device_class_file")
+    if [ "$device_class" != "0x030000" ]; then
+      echo "NONE"
+      return
+    fi
+
+    if ls "/sys/bus/pci/devices/$gpu" | grep consumer >& /dev/null; then
+      aux_dev=$(ls "/sys/bus/pci/devices/$gpu" | grep consumer | awk -Fconsumer:pci: '{print $2}')
+      if [ "$aux_dev" == "" ]; then
+        echo "NONE"
+        return
+      fi
+
+      if ls "/sys/bus/pci/devices/$aux_dev/" >& /dev/null; then
+        echo "$aux_dev"
+        return
+      fi
+    fi
+
+    echo "NONE"
+}
+
+bind_device() {
+    local gpu=$1
+
+    if ! is_nvidia_gpu_device $gpu; then
+        echo "device $gpu is not a gpu!"
+        return 0
+    fi
+
+    bind_pci_device "$gpu"
+    #for graphics mode, we need to bind the auxiliary device as well
+    aux_dev=$(get_graphics_aux_dev "$gpu")
+    if [ "$aux_dev" != "NONE" ]; then
+      echo "gpu $gpu is in graphics mode aux_dev $aux_dev"
+      bind_pci_device "$aux_dev"
+    fi
+}
+
 bind_all() {
     for dev in /sys/bus/pci/devices/*; do
         read vendor < $dev/vendor