"NPU_OCAPI" };
 
 static void pnv_pci_ioda2_set_bypass(struct pnv_ioda_pe *pe, bool enable);
+static void pnv_pci_configure_bus(struct pci_bus *bus);
 
 void pe_level_printk(const struct pnv_ioda_pe *pe, const char *level,
                            const char *fmt, ...)
        return pe;
 }
 
-static void pnv_ioda_setup_same_PE(struct pci_bus *bus, struct pnv_ioda_pe *pe)
-{
-       struct pci_dev *dev;
-
-       list_for_each_entry(dev, &bus->devices, bus_list) {
-               struct pci_dn *pdn = pci_get_pdn(dev);
-
-               if (pdn == NULL) {
-                       pr_warn("%s: No device node associated with device !\n",
-                               pci_name(dev));
-                       continue;
-               }
-
-               /*
-                * In partial hotplug case, the PCI device might be still
-                * associated with the PE and needn't attach it to the PE
-                * again.
-                */
-               if (pdn->pe_number != IODA_INVALID_PE)
-                       continue;
-
-               pe->device_count++;
-               pdn->pe_number = pe->pe_number;
-               if ((pe->flags & PNV_IODA_PE_BUS_ALL) && dev->subordinate)
-                       pnv_ioda_setup_same_PE(dev->subordinate, pe);
-       }
-}
-
 /*
  * There're 2 types of PCI bus sensitive PEs: One that is compromised of
  * single PCI bus. Another one that contains the primary PCI bus and its
        pe_num = phb->ioda.pe_rmap[bus->number << 8];
        if (pe_num != IODA_INVALID_PE) {
                pe = &phb->ioda.pe_array[pe_num];
-               pnv_ioda_setup_same_PE(bus, pe);
                return NULL;
        }
 
                return NULL;
        }
 
-       /* Associate it with all child devices */
-       pnv_ioda_setup_same_PE(bus, pe);
-
        /* Put PE to the list */
        list_add_tail(&pe->list, &phb->ioda.pe_list);
 
        struct pci_dn *pdn = pci_get_pdn(pdev);
        struct pnv_ioda_pe *pe;
 
-       /*
-        * The function can be called while the PE#
-        * hasn't been assigned. Do nothing for the
-        * case.
-        */
-       if (!pdn || pdn->pe_number == IODA_INVALID_PE)
-               return;
+       /* Check if the BDFN for this device is associated with a PE yet */
+       pe = pnv_pci_bdfn_to_pe(phb, pdev->devfn | (pdev->bus->number << 8));
+       if (!pe) {
+               /* VF PEs should be pre-configured in pnv_pci_sriov_enable() */
+               if (WARN_ON(pdev->is_virtfn))
+                       return;
+
+               pnv_pci_configure_bus(pdev->bus);
+               pe = pnv_pci_bdfn_to_pe(phb, pdev->devfn | (pdev->bus->number << 8));
+               pci_info(pdev, "Configured PE#%x\n", pe ? pe->pe_number : 0xfffff);
+
+
+               /*
+                * If we can't setup the IODA PE something has gone horribly
+                * wrong and we can't enable DMA for the device.
+                */
+               if (WARN_ON(!pe))
+                       return;
+       } else {
+               pci_info(pdev, "Added to existing PE#%x\n", pe->pe_number);
+       }
+
+       if (pdn)
+               pdn->pe_number = pe->pe_number;
+       pe->device_count++;
 
-       pe = &phb->ioda.pe_array[pdn->pe_number];
        WARN_ON(get_dma_ops(&pdev->dev) != &dma_iommu_ops);
        pdev->dev.archdata.dma_offset = pe->tce_bypass_base;
        set_iommu_table_base(&pdev->dev, pe->table_group.tables[0]);
        pe->table_group.tce32_size = tbl->it_size << tbl->it_page_shift;
        iommu_init_table(tbl, phb->hose->node, 0, 0);
 
-       if (pe->flags & (PNV_IODA_PE_BUS | PNV_IODA_PE_BUS_ALL))
-               pnv_ioda_setup_bus_dma(pe, pe->pbus);
-
        return;
  fail:
        /* XXX Failure: Try to fallback to 64-bit only ? */
        iommu_register_group(&pe->table_group, phb->hose->global_number,
                             pe->pe_number);
 #endif
-
-       if (pe->flags & (PNV_IODA_PE_BUS | PNV_IODA_PE_BUS_ALL))
-               pnv_ioda_setup_bus_dma(pe, pe->pbus);
 }
 
 int64_t pnv_opal_pci_msi_eoi(struct irq_chip *chip, unsigned int hw_irq)
        }
 }
 
-static void pnv_pci_setup_bridge(struct pci_bus *bus, unsigned long type)
+static void pnv_pci_configure_bus(struct pci_bus *bus)
 {
        struct pci_controller *hose = pci_bus_to_host(bus);
        struct pnv_phb *phb = hose->private_data;
        struct pci_dev *bridge = bus->self;
        struct pnv_ioda_pe *pe;
-       bool all = (pci_pcie_type(bridge) == PCI_EXP_TYPE_PCI_BRIDGE);
+       bool all = (bridge && pci_pcie_type(bridge) == PCI_EXP_TYPE_PCI_BRIDGE);
 
-       /* Extend bridge's windows if necessary */
-       pnv_pci_fixup_bridge_resources(bus, type);
+       dev_info(&bus->dev, "Configuring PE for bus\n");
 
        /* The PE for root bus should be realized before any one else */
        if (!phb->ioda.root_pe_populated) {
        .enable_device_hook     = pnv_pci_enable_device_hook,
        .release_device         = pnv_pci_release_device,
        .window_alignment       = pnv_pci_window_alignment,
-       .setup_bridge           = pnv_pci_setup_bridge,
+       .setup_bridge           = pnv_pci_fixup_bridge_resources,
        .reset_secondary_bus    = pnv_pci_reset_secondary_bus,
        .shutdown               = pnv_pci_ioda_shutdown,
 };