return 0;
 }
 
+#define ISST_MAX_BUS_NUMBER    2
 
 struct isst_if_cpu_info {
        /* For BUS 0 and BUS 1 only, which we need for PUNIT interface */
-       int bus_info[2];
-       struct pci_dev *pci_dev[2];
+       int bus_info[ISST_MAX_BUS_NUMBER];
+       struct pci_dev *pci_dev[ISST_MAX_BUS_NUMBER];
        int punit_cpu_id;
        int numa_node;
 };
 
+struct isst_if_pkg_info {
+       struct pci_dev *pci_dev[ISST_MAX_BUS_NUMBER];
+};
+
 static struct isst_if_cpu_info *isst_cpu_info;
+static struct isst_if_pkg_info *isst_pkg_info;
+
 #define ISST_MAX_PCI_DOMAINS   8
 
 static struct pci_dev *_isst_if_get_pci_dev(int cpu, int bus_no, int dev, int fn)
 {
        struct pci_dev *matched_pci_dev = NULL;
        struct pci_dev *pci_dev = NULL;
-       int no_matches = 0;
+       int no_matches = 0, pkg_id;
        int i, bus_number;
 
-       if (bus_no < 0 || bus_no > 1 || cpu < 0 || cpu >= nr_cpu_ids ||
-           cpu >= num_possible_cpus())
+       if (bus_no < 0 || bus_no >= ISST_MAX_BUS_NUMBER || cpu < 0 ||
+           cpu >= nr_cpu_ids || cpu >= num_possible_cpus())
                return NULL;
 
+       pkg_id = topology_physical_package_id(cpu);
+
        bus_number = isst_cpu_info[cpu].bus_info[bus_no];
        if (bus_number < 0)
                return NULL;
                }
 
                if (node == isst_cpu_info[cpu].numa_node) {
+                       isst_pkg_info[pkg_id].pci_dev[bus_no] = _pci_dev;
+
                        pci_dev = _pci_dev;
                        break;
                }
        if (!pci_dev && no_matches == 1)
                pci_dev = matched_pci_dev;
 
+       /* Return pci_dev pointer for any matched CPU in the package */
+       if (!pci_dev)
+               pci_dev = isst_pkg_info[pkg_id].pci_dev[bus_no];
+
        return pci_dev;
 }
 
 {
        struct pci_dev *pci_dev;
 
-       if (bus_no < 0 || bus_no > 1 || cpu < 0 || cpu >= nr_cpu_ids ||
-           cpu >= num_possible_cpus())
+       if (bus_no < 0 || bus_no >= ISST_MAX_BUS_NUMBER  || cpu < 0 ||
+           cpu >= nr_cpu_ids || cpu >= num_possible_cpus())
                return NULL;
 
        pci_dev = isst_cpu_info[cpu].pci_dev[bus_no];
        if (!isst_cpu_info)
                return -ENOMEM;
 
+       isst_pkg_info = kcalloc(topology_max_packages(),
+                               sizeof(*isst_pkg_info),
+                               GFP_KERNEL);
+       if (!isst_pkg_info) {
+               kfree(isst_cpu_info);
+               return -ENOMEM;
+       }
+
        ret = cpuhp_setup_state(CPUHP_AP_ONLINE_DYN,
                                "platform/x86/isst-if:online",
                                isst_if_cpu_online, NULL);
        if (ret < 0) {
+               kfree(isst_pkg_info);
                kfree(isst_cpu_info);
                return ret;
        }
 static void isst_if_cpu_info_exit(void)
 {
        cpuhp_remove_state(isst_if_online_id);
+       kfree(isst_pkg_info);
        kfree(isst_cpu_info);
 };