From 66f53fa89ea2ada5e7bdf956dc35c47789ddad3b Mon Sep 17 00:00:00 2001
From: Qiu Jian <qiujian@yunionyun.com>
Date: Mon, 13 Nov 2023 16:19:02 +0800
Subject: [PATCH] fix: extend gpu pci class codes

---
 pkg/baremetal/tasks/baseprepare.go |  6 ++---
 pkg/hostman/isolated_device/gpu.go | 43 +++++++++++-------------------
 2 files changed, 18 insertions(+), 31 deletions(-)

diff --git a/pkg/baremetal/tasks/baseprepare.go b/pkg/baremetal/tasks/baseprepare.go
index dbf407d6bd..59f0afd3e5 100644
--- a/pkg/baremetal/tasks/baseprepare.go
+++ b/pkg/baremetal/tasks/baseprepare.go
@@ -735,9 +735,9 @@ func getIsolatedDevicesInfo(cli *ssh.Client, ip net.IP) ([]*isolated_device.PCID
 		bootVgaPath = append(bootVgaPath, strings.TrimSpace(lines[i]))
 	}
 
-	cmd := "lspci -nnmm | egrep '3D|VGA'"
+	cmd := "lspci -nnmm"
 	if updatedPciids {
-		cmd = "lspci -i /pci.ids -nnmm | egrep '3D|VGA'"
+		cmd = "lspci -i /pci.ids -nnmm"
 	}
 
 	lines, err = cli.Run(cmd)
@@ -748,7 +748,7 @@ func getIsolatedDevicesInfo(cli *ssh.Client, ip net.IP) ([]*isolated_device.PCID
 	for _, line := range lines {
 		if len(line) > 0 {
 			dev := isolated_device.NewPCIDevice2(line)
-			if len(dev.Addr) > 0 && !isBootVga(cli, dev, bootVgaPath) {
+			if len(dev.Addr) > 0 && utils.IsInArray(dev.ClassCode, isolated_device.GpuClassCodes) && !isBootVga(cli, dev, bootVgaPath) {
 				devs = append(devs, dev)
 			}
 		}
diff --git a/pkg/hostman/isolated_device/gpu.go b/pkg/hostman/isolated_device/gpu.go
index ad4207fcfe..1064ada6b4 100644
--- a/pkg/hostman/isolated_device/gpu.go
+++ b/pkg/hostman/isolated_device/gpu.go
@@ -39,6 +39,16 @@ import (
 const (
 	CLASS_CODE_VGA = "0300"
 	CLASS_CODE_3D  = "0302"
+
+	CLASS_CODE_DISP = "0380"
+)
+
+var (
+	GpuClassCodes = []string{
+		CLASS_CODE_VGA,
+		CLASS_CODE_3D,
+		CLASS_CODE_DISP,
+	}
 )
 
 const (
@@ -69,6 +79,9 @@ func getPassthroughGPUS(filteredAddrs []string) ([]*PCIDevice, error, []error) {
 		if utils.IsInStringArray(dev.Addr, filteredAddrs) {
 			continue
 		}
+		if !utils.IsInArray(dev.ClassCode, GpuClassCodes) {
+			continue
+		}
 		if err := dev.checkSameIOMMUGroupDevice(); err != nil {
 			warns = append(warns, errors.Wrapf(err, "get dev %s iommu group devices", dev.Addr))
 			continue
@@ -96,7 +109,7 @@ func getPassthroughGPUS(filteredAddrs []string) ([]*PCIDevice, error, []error) {
 }
 
 func getGPUPCIStr() ([]string, error) {
-	cmd := "lspci -nnmm | egrep '3D|VGA'"
+	cmd := "lspci -nnmm"
 	ret, err := bashOutput(cmd)
 	if err != nil {
 		return nil, err
@@ -236,32 +249,6 @@ func (gpu *sGPUHPCDevice) GetPassthroughCmd(index int) string {
 	return fmt.Sprintf(" -device vfio-pci,host=%s,multifunction=on", gpu.GetAddr())
 }
 
-func gpuPCIString() ([]string, error) {
-	lines, err := bashOutput("lspci -nnmm | egrep '3D|VGA'")
-	if err != nil {
-		return nil, fmt.Errorf("Get GPU PCI: %v", err)
-	}
-	ret := []string{}
-	for _, line := range lines {
-		if len(line) != 0 {
-			ret = append(ret, line)
-		}
-	}
-	return ret, nil
-}
-
-func gpuPCIAddr() ([]string, error) {
-	lines, err := gpuPCIString()
-	if err != nil {
-		return nil, err
-	}
-	addrs := []string{}
-	for _, line := range lines {
-		addrs = append(addrs, strings.Split(line, " ")[0])
-	}
-	return addrs, nil
-}
-
 // parseLspci parse one line output of `lspci -nnmm`
 func parseLspci(line string) *PCIDevice {
 	itemRegex := `(?P<bus_id>(` + BUSID_REGEX + `))` +
@@ -329,7 +316,7 @@ func (d *PCIDevice) IsBootVGA() (bool, error) {
 }
 
 func (d *PCIDevice) forceBindVFIOPCIDriver(useBootVGA bool) error {
-	if !utils.IsInStringArray(d.ClassCode, []string{CLASS_CODE_VGA, CLASS_CODE_3D}) {
+	if !utils.IsInArray(d.ClassCode, GpuClassCodes) {
 		return nil
 	}
 	isBootVGA, err := d.IsBootVGA()