4 files changed, 17 insertions, 19 deletions
diff --git a/payloads/libpayload/drivers/storage/ahci.c b/payloads/libpayload/drivers/storage/ahci.c
index ff22b19b71..e54bfbfef0 100644
--- a/payloads/libpayload/drivers/storage/ahci.c
+++ b/payloads/libpayload/drivers/storage/ahci.c
@@ -227,7 +227,8 @@ static u32 working_controllers[] = {
 	0x8086 | 0x5ae3 << 16, /* Apollo Lake */
 };
 #endif
-static void ahci_init_pci(pcidev_t dev)
+
+void ahci_initialize(pcidev_t dev)
 {
 	int i;
 
@@ -282,19 +283,3 @@ static void ahci_init_pci(pcidev_t dev)
 			ahci_port_probe(ctrl, &ports[i], i + 1);
 	}
 }
-
-void ahci_initialize(void)
-{
-	int bus, dev, func;
-
-	for (bus = 0; bus < 256; ++bus) {
-		for (dev = 0; dev < 32; ++dev) {
-			const u16 class =
-				pci_read_config16(PCI_DEV(bus, dev, 0), 0xa);
-			if (class != 0xffff) {
-				for (func = 0; func < 8; ++func)
-					ahci_init_pci(PCI_DEV(bus, dev, func));
-			}
-		}
-	}
-}
diff --git a/payloads/libpayload/drivers/storage/storage.c b/payloads/libpayload/drivers/storage/storage.c
index 24b9024a93..addf8d2bd3 100644
--- a/payloads/libpayload/drivers/storage/storage.c
+++ b/payloads/libpayload/drivers/storage/storage.c
@@ -27,6 +27,7 @@
  */
 
 #include <libpayload.h>
+#include <pci/pci.h>
 #if CONFIG(LP_STORAGE_AHCI)
 # include <storage/ahci.h>
 #endif
@@ -108,7 +109,18 @@ ssize_t storage_read_blocks512(const size_t dev_num,
  */
 void storage_initialize(void)
 {
+#if CONFIG(LP_PCI)
+	struct pci_dev *dev;
+	for (dev = lib_sysinfo.pacc.devices; dev; dev = dev->next) {
+		switch (dev->device_class) {
 #if CONFIG(LP_STORAGE_AHCI)
-	ahci_initialize();
+		case PCI_CLASS_STORAGE_AHCI:
+			ahci_initialize(PCI_DEV(dev->bus, dev->dev, dev->func));
+			break;
+#endif
+		default:
+			break;
+		}
+	}
 #endif
 }
diff --git a/payloads/libpayload/include/pci/pci.h b/payloads/libpayload/include/pci/pci.h
index 7ec9491d95..5e21060130 100644
--- a/payloads/libpayload/include/pci/pci.h
+++ b/payloads/libpayload/include/pci/pci.h
@@ -66,6 +66,7 @@
 #define PCI_ROM_ADDRESS1	0x38 // on bridges
 #define PCI_ROM_ADDRESS_MASK	~0x7ff
 
+#define PCI_CLASS_STORAGE_AHCI	0x0106
 #define PCI_CLASS_MEMORY_OTHER	0x0580
 
 #define PCI_VENDOR_ID_INTEL 0x8086
diff --git a/payloads/libpayload/include/storage/ahci.h b/payloads/libpayload/include/storage/ahci.h
index 0aa032297a..28436249f8 100644
--- a/payloads/libpayload/include/storage/ahci.h
+++ b/payloads/libpayload/include/storage/ahci.h
@@ -29,6 +29,6 @@
 #ifndef _STORAGE_AHCI_H
 #define _STORAGE_AHCI_H
 
-void ahci_initialize(void);
+void ahci_initialize(pcidev_t dev);
 
 #endif