u16 length;
        u16 devid;
        u16 aux;
-       u64 resv;
+       u16 pci_seg;
+       u8  resv[6];
        u64 range_start;
        u64 range_length;
 } __attribute__((packed));
 
 u16 amd_iommu_last_bdf;                        /* largest PCI device id we have
                                           to handle */
-LIST_HEAD(amd_iommu_unity_map);                /* a list of required unity mappings
-                                          we find in ACPI */
 
 LIST_HEAD(amd_iommu_pci_seg_list);     /* list of all PCI segments */
 LIST_HEAD(amd_iommu_list);             /* list of all AMD IOMMUs in the
 
        pci_seg->id = id;
        init_llist_head(&pci_seg->dev_data_list);
+       INIT_LIST_HEAD(&pci_seg->unity_map);
        list_add_tail(&pci_seg->list, &amd_iommu_pci_seg_list);
 
        if (alloc_dev_table(pci_seg))
 static void __init free_unity_maps(void)
 {
        struct unity_map_entry *entry, *next;
+       struct amd_iommu_pci_seg *p, *pci_seg;
 
-       list_for_each_entry_safe(entry, next, &amd_iommu_unity_map, list) {
-               list_del(&entry->list);
-               kfree(entry);
+       for_each_pci_segment_safe(pci_seg, p) {
+               list_for_each_entry_safe(entry, next, &pci_seg->unity_map, list) {
+                       list_del(&entry->list);
+                       kfree(entry);
+               }
        }
 }
 
 static int __init init_unity_map_range(struct ivmd_header *m)
 {
        struct unity_map_entry *e = NULL;
+       struct amd_iommu_pci_seg *pci_seg;
        char *s;
 
+       pci_seg = get_pci_segment(m->pci_seg);
+       if (pci_seg == NULL)
+               return -ENOMEM;
+
        e = kzalloc(sizeof(*e), GFP_KERNEL);
        if (e == NULL)
                return -ENOMEM;
        if (m->flags & IVMD_FLAG_EXCL_RANGE)
                e->prot = (IVMD_FLAG_IW | IVMD_FLAG_IR) >> 1;
 
-       DUMP_printk("%s devid_start: %02x:%02x.%x devid_end: %02x:%02x.%x"
-                   " range_start: %016llx range_end: %016llx flags: %x\n", s,
+       DUMP_printk("%s devid_start: %04x:%02x:%02x.%x devid_end: "
+                   "%04x:%02x:%02x.%x range_start: %016llx range_end: %016llx"
+                   " flags: %x\n", s, m->pci_seg,
                    PCI_BUS_NUM(e->devid_start), PCI_SLOT(e->devid_start),
-                   PCI_FUNC(e->devid_start), PCI_BUS_NUM(e->devid_end),
+                   PCI_FUNC(e->devid_start), m->pci_seg,
+                   PCI_BUS_NUM(e->devid_end),
                    PCI_SLOT(e->devid_end), PCI_FUNC(e->devid_end),
                    e->address_start, e->address_end, m->flags);
 
-       list_add_tail(&e->list, &amd_iommu_unity_map);
+       list_add_tail(&e->list, &pci_seg->unity_map);
 
        return 0;
 }
 
 {
        struct iommu_resv_region *region;
        struct unity_map_entry *entry;
+       struct amd_iommu *iommu;
+       struct amd_iommu_pci_seg *pci_seg;
        int devid;
 
        devid = get_device_id(dev);
        if (devid < 0)
                return;
+       iommu = rlookup_amd_iommu(dev);
+       if (!iommu)
+               return;
+       pci_seg = iommu->pci_seg;
 
-       list_for_each_entry(entry, &amd_iommu_unity_map, list) {
+       list_for_each_entry(entry, &pci_seg->unity_map, list) {
                int type, prot = 0;
                size_t length;