static int amd_iommu_enable_interrupts(void);
 static int __init iommu_go_to_state(enum iommu_init_state state);
-static void init_device_table_dma(void);
+static void init_device_table_dma(struct amd_iommu_pci_seg *pci_seg);
 
 static bool amd_iommu_pre_enabled = true;
 
 static int __init amd_iommu_init_pci(void)
 {
        struct amd_iommu *iommu;
+       struct amd_iommu_pci_seg *pci_seg;
        int ret;
 
        for_each_iommu(iommu) {
                goto out;
        }
 
-       init_device_table_dma();
+       for_each_pci_segment(pci_seg)
+               init_device_table_dma(pci_seg);
 
        for_each_iommu(iommu)
                iommu_flush_all_caches(iommu);
 /*
  * Init the device table to not allow DMA access for devices
  */
-static void init_device_table_dma(void)
+static void init_device_table_dma(struct amd_iommu_pci_seg *pci_seg)
 {
        u32 devid;
+       struct dev_table_entry *dev_table = pci_seg->dev_table;
+
+       if (dev_table == NULL)
+               return;
 
        for (devid = 0; devid <= amd_iommu_last_bdf; ++devid) {
                set_dev_entry_bit(devid, DEV_ENTRY_VALID);
        }
 }
 
-static void __init uninit_device_table_dma(void)
+static void __init uninit_device_table_dma(struct amd_iommu_pci_seg *pci_seg)
 {
        u32 devid;
+       struct dev_table_entry *dev_table = pci_seg->dev_table;
+
+       if (dev_table == NULL)
+               return;
 
        for (devid = 0; devid <= amd_iommu_last_bdf; ++devid) {
-               amd_iommu_dev_table[devid].data[0] = 0ULL;
-               amd_iommu_dev_table[devid].data[1] = 0ULL;
+               dev_table[devid].data[0] = 0ULL;
+               dev_table[devid].data[1] = 0ULL;
        }
 }
 
                        free_iommu_resources();
                } else {
                        struct amd_iommu *iommu;
+                       struct amd_iommu_pci_seg *pci_seg;
+
+                       for_each_pci_segment(pci_seg)
+                               uninit_device_table_dma(pci_seg);
 
-                       uninit_device_table_dma();
                        for_each_iommu(iommu)
                                iommu_flush_all_caches(iommu);
                }