extern void iommu_register_group(struct iommu_table_group *table_group,
                                 int pci_domain_number, unsigned long pe_num);
-extern int iommu_add_device(struct device *dev);
+extern int iommu_add_device(struct iommu_table_group *table_group,
+               struct device *dev);
 extern void iommu_del_device(struct device *dev);
-extern int __init tce_iommu_bus_notifier_init(void);
 extern long iommu_tce_xchg(struct mm_struct *mm, struct iommu_table *tbl,
                unsigned long entry, unsigned long *hpa,
                enum dma_data_direction *direction);
 {
 }
 
-static inline int iommu_add_device(struct device *dev)
+static inline int iommu_add_device(struct iommu_table_group *table_group,
+               struct device *dev)
 {
        return 0;
 }
 static inline void iommu_del_device(struct device *dev)
 {
 }
-
-static inline int __init tce_iommu_bus_notifier_init(void)
-{
-        return 0;
-}
 #endif /* !CONFIG_IOMMU_API */
 
 int dma_iommu_mapping_error(struct device *dev, dma_addr_t dma_addr);
 
 }
 EXPORT_SYMBOL_GPL(iommu_release_ownership);
 
-int iommu_add_device(struct device *dev)
+int iommu_add_device(struct iommu_table_group *table_group, struct device *dev)
 {
-       struct iommu_table *tbl;
-       struct iommu_table_group_link *tgl;
-
        /*
         * The sysfs entries should be populated before
         * binding IOMMU group. If sysfs entries isn't
                return -EBUSY;
        }
 
-       tbl = get_iommu_table_base(dev);
-       if (!tbl) {
-               pr_debug("%s: Skipping device %s with no tbl\n",
-                        __func__, dev_name(dev));
-               return 0;
-       }
-
-       tgl = list_first_entry_or_null(&tbl->it_group_list,
-                       struct iommu_table_group_link, next);
-       if (!tgl) {
-               pr_debug("%s: Skipping device %s with no group\n",
-                        __func__, dev_name(dev));
-               return 0;
-       }
        pr_debug("%s: Adding %s to iommu group %d\n",
-                __func__, dev_name(dev),
-                iommu_group_id(tgl->table_group->group));
-
-       if (PAGE_SIZE < IOMMU_PAGE_SIZE(tbl)) {
-               pr_err("%s: Invalid IOMMU page size %lx (%lx) on %s\n",
-                      __func__, IOMMU_PAGE_SIZE(tbl),
-                      PAGE_SIZE, dev_name(dev));
-               return -EINVAL;
-       }
+                __func__, dev_name(dev),  iommu_group_id(table_group->group));
 
-       return iommu_group_add_device(tgl->table_group->group, dev);
+       return iommu_group_add_device(table_group->group, dev);
 }
 EXPORT_SYMBOL_GPL(iommu_add_device);
 
        iommu_group_remove_device(dev);
 }
 EXPORT_SYMBOL_GPL(iommu_del_device);
-
-static int tce_iommu_bus_notifier(struct notifier_block *nb,
-                unsigned long action, void *data)
-{
-        struct device *dev = data;
-
-        switch (action) {
-        case BUS_NOTIFY_ADD_DEVICE:
-                return iommu_add_device(dev);
-        case BUS_NOTIFY_DEL_DEVICE:
-                if (dev->iommu_group)
-                        iommu_del_device(dev);
-                return 0;
-        default:
-                return 0;
-        }
-}
-
-static struct notifier_block tce_iommu_bus_nb = {
-        .notifier_call = tce_iommu_bus_notifier,
-};
-
-int __init tce_iommu_bus_notifier_init(void)
-{
-        bus_register_notifier(&pci_bus_type, &tce_iommu_bus_nb);
-        return 0;
-}
 #endif /* CONFIG_IOMMU_API */
 
                set_iommu_table_base(&dev->dev, pe->table_group.tables[0]);
                set_dma_offset(&dev->dev, pe->tce_bypass_base);
                if (add_to_group)
-                       iommu_add_device(&dev->dev);
+                       iommu_add_device(&pe->table_group, &dev->dev);
 
                if ((pe->flags & PNV_IODA_PE_BUS_ALL) && dev->subordinate)
                        pnv_ioda_setup_bus_dma(pe, dev->subordinate,
        if (!pnv_iommu_bypass_disabled)
                pnv_pci_ioda2_set_bypass(pe, true);
 
-       /*
-        * Setting table base here only for carrying iommu_group
-        * further down to let iommu_add_device() do the job.
-        * pnv_pci_ioda_dma_dev_setup will override it later anyway.
-        */
-       if (pe->flags & PNV_IODA_PE_DEV)
-               set_iommu_table_base(&pe->pdev->dev, tbl);
-
        return 0;
 }
 
 
        set_pci_dma_ops(&dma_iommu_ops);
 }
 
-machine_subsys_initcall_sync(powernv, tce_iommu_bus_notifier_init);
+static int pnv_tce_iommu_bus_notifier(struct notifier_block *nb,
+               unsigned long action, void *data)
+{
+       struct device *dev = data;
+       struct pci_dev *pdev;
+       struct pci_dn *pdn;
+       struct pnv_ioda_pe *pe;
+       struct pci_controller *hose;
+       struct pnv_phb *phb;
+
+       switch (action) {
+       case BUS_NOTIFY_ADD_DEVICE:
+               pdev = to_pci_dev(dev);
+               pdn = pci_get_pdn(pdev);
+               hose = pci_bus_to_host(pdev->bus);
+               phb = hose->private_data;
+
+               WARN_ON_ONCE(!phb);
+               if (!pdn || pdn->pe_number == IODA_INVALID_PE || !phb)
+                       return 0;
+
+               pe = &phb->ioda.pe_array[pdn->pe_number];
+               iommu_add_device(&pe->table_group, dev);
+               return 0;
+       case BUS_NOTIFY_DEL_DEVICE:
+               iommu_del_device(dev);
+               return 0;
+       default:
+               return 0;
+       }
+}
+
+static struct notifier_block pnv_tce_iommu_bus_nb = {
+       .notifier_call = pnv_tce_iommu_bus_notifier,
+};
+
+static int __init pnv_tce_iommu_bus_notifier_init(void)
+{
+       bus_register_notifier(&pci_bus_type, &pnv_tce_iommu_bus_nb);
+       return 0;
+}
+machine_subsys_initcall_sync(powernv, pnv_tce_iommu_bus_notifier_init);
 
 {
        struct iommu_table_group *table_group;
        struct iommu_table *tbl;
-       struct iommu_table_group_link *tgl;
 
        table_group = kzalloc_node(sizeof(struct iommu_table_group), GFP_KERNEL,
                           node);
        if (!tbl)
                goto free_group;
 
-       tgl = kzalloc_node(sizeof(struct iommu_table_group_link), GFP_KERNEL,
-                       node);
-       if (!tgl)
-               goto free_table;
-
        INIT_LIST_HEAD_RCU(&tbl->it_group_list);
        kref_init(&tbl->it_kref);
-       tgl->table_group = table_group;
-       list_add_rcu(&tgl->next, &tbl->it_group_list);
 
        table_group->tables[0] = tbl;
 
        return table_group;
 
-free_table:
-       kfree(tbl);
 free_group:
        kfree(table_group);
        return NULL;
                const char *node_name)
 {
        struct iommu_table *tbl;
-#ifdef CONFIG_IOMMU_API
-       struct iommu_table_group_link *tgl;
-#endif
 
        if (!table_group)
                return;
 
        tbl = table_group->tables[0];
 #ifdef CONFIG_IOMMU_API
-       tgl = list_first_entry_or_null(&tbl->it_group_list,
-                       struct iommu_table_group_link, next);
-
-       WARN_ON_ONCE(!tgl);
-       if (tgl) {
-               list_del_rcu(&tgl->next);
-               kfree(tgl);
-       }
        if (table_group->group) {
                iommu_group_put(table_group->group);
                BUG_ON(table_group->group);
        }
 
        set_iommu_table_base(&dev->dev, pci->table_group->tables[0]);
-       iommu_add_device(&dev->dev);
+       iommu_add_device(pci->table_group, &dev->dev);
 }
 
 static int dma_set_mask_pSeriesLP(struct device *dev, u64 dma_mask)
 
 __setup("multitce=", disable_multitce);
 
+static int tce_iommu_bus_notifier(struct notifier_block *nb,
+               unsigned long action, void *data)
+{
+       struct device *dev = data;
+
+       switch (action) {
+       case BUS_NOTIFY_DEL_DEVICE:
+               iommu_del_device(dev);
+               return 0;
+       default:
+               return 0;
+       }
+}
+
+static struct notifier_block tce_iommu_bus_nb = {
+       .notifier_call = tce_iommu_bus_notifier,
+};
+
+static int __init tce_iommu_bus_notifier_init(void)
+{
+       bus_register_notifier(&pci_bus_type, &tce_iommu_bus_nb);
+       return 0;
+}
 machine_subsys_initcall_sync(pseries, tce_iommu_bus_notifier_init);