}
 }
 
-static void pnv_pci_ioda_dma_dev_setup(struct pnv_phb *phb, struct pci_dev *dev)
+static void pnv_pci_ioda_dma_dev_setup(struct pnv_phb *phb, struct pci_dev *pdev)
 {
-       /* We delay DMA setup after we have assigned all PE# */
-}
+       struct pci_dn *pdn = pnv_ioda_get_pdn(pdev);
+       struct pnv_ioda_pe *pe;
 
-static void pnv_ioda_setup_bus_dma(struct pnv_ioda_pe *pe, struct pci_bus *bus)
-{
-       struct pci_dev *dev;
+       /*
+        * 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;
 
-       list_for_each_entry(dev, &bus->devices, bus_list) {
-               set_iommu_table_base(&dev->dev, &pe->tce32_table);
-               if (dev->subordinate)
-                       pnv_ioda_setup_bus_dma(pe, dev->subordinate);
-       }
+       pe = &phb->ioda.pe_array[pdn->pe_number];
+       set_iommu_table_base(&pdev->dev, &pe->tce32_table);
 }
 
 static void pnv_pci_ioda1_tce_invalidate(struct iommu_table *tbl,
        }
        iommu_init_table(tbl, phb->hose->node);
 
-       if (pe->pdev)
-               set_iommu_table_base(&pe->pdev->dev, tbl);
-       else
-               pnv_ioda_setup_bus_dma(pe, pe->pbus);
-
        return;
  fail:
        /* XXX Failure: Try to fallback to 64-bit only ? */
        }
        iommu_init_table(tbl, phb->hose->node);
 
-       if (pe->pdev)
-               set_iommu_table_base(&pe->pdev->dev, tbl);
-       else
-               pnv_ioda_setup_bus_dma(pe, pe->pbus);
-
        return;
 fail:
        if (pe->tce32_seg >= 0)