diff --git a/xen/drivers/passthrough/amd/iommu_cmd.c b/xen/drivers/passthrough/amd/iommu_cmd.c
index 7c9d9be..b3094f3 100644
--- a/xen/drivers/passthrough/amd/iommu_cmd.c
+++ b/xen/drivers/passthrough/amd/iommu_cmd.c
@@ -298,24 +298,23 @@ void amd_iommu_flush_iotlb(u8 devfn, const struct pci_dev *pdev,
     if ( ats_pdev == NULL )
         return;
 
-    if ( !pci_ats_enabled(ats_pdev->seg, ats_pdev->bus, ats_pdev->devfn) )
+    if ( !pci_ats_enabled(pdev->seg, pdev->bus, pdev->devfn) )
         return;
 
-    iommu = find_iommu_for_device(ats_pdev->seg,
-                                  PCI_BDF2(ats_pdev->bus, ats_pdev->devfn));
+    iommu = find_iommu_for_device(pdev->seg, PCI_BDF2(pdev->bus, pdev->devfn));
 
     if ( !iommu )
     {
         AMD_IOMMU_DEBUG("%s: Can't find iommu for %04x:%02x:%02x.%u\n",
-                        __func__, ats_pdev->seg, ats_pdev->bus,
-                        PCI_SLOT(ats_pdev->devfn), PCI_FUNC(ats_pdev->devfn));
+                        __func__, pdev->seg, pdev->bus,
+                        PCI_SLOT(pdev->devfn), PCI_FUNC(pdev->devfn));
         return;
     }
 
     if ( !iommu_has_cap(iommu, PCI_CAP_IOTLB_SHIFT) )
         return;
 
-    req_id = get_dma_requestor_id(iommu->seg, PCI_BDF2(ats_pdev->bus, devfn));
+    req_id = get_dma_requestor_id(iommu->seg, PCI_BDF2(pdev->bus, devfn));
     queueid = req_id;
     maxpend = ats_pdev->ats_queue_depth & 0xff;
 
diff --git a/xen/drivers/passthrough/amd/pci_amd_iommu.c b/xen/drivers/passthrough/amd/pci_amd_iommu.c
index 7761241..64ca78e 100644
--- a/xen/drivers/passthrough/amd/pci_amd_iommu.c
+++ b/xen/drivers/passthrough/amd/pci_amd_iommu.c
@@ -162,7 +162,7 @@ static void amd_iommu_setup_domain_device(
          !pci_ats_enabled(iommu->seg, bus, pdev->devfn) )
     {
         if ( devfn == pdev->devfn )
-            enable_ats_device(iommu->seg, bus, devfn, iommu);
+            enable_ats_device(iommu, pdev);
 
         amd_iommu_flush_iotlb(devfn, pdev, INV_IOMMU_ALL_PAGES_ADDRESS, 0);
     }
diff --git a/xen/drivers/passthrough/ats.h b/xen/drivers/passthrough/ats.h
index 5c91572..1359841 100644
--- a/xen/drivers/passthrough/ats.h
+++ b/xen/drivers/passthrough/ats.h
@@ -19,9 +19,7 @@
 
 struct pci_ats_dev {
     struct list_head list;
-    u16 seg;
-    u8 bus;
-    u8 devfn;
+    struct pci_dev *pci_dev;
     u16 ats_queue_depth;    /* ATS device invalidation queue depth */
     const void *iommu;      /* No common IOMMU struct so use void pointer */
 };
@@ -34,7 +32,7 @@ struct pci_ats_dev {
 extern struct list_head ats_devices;
 extern bool_t ats_enabled;
 
-int enable_ats_device(int seg, int bus, int devfn, const void *iommu);
+int enable_ats_device(const void *iommu, struct pci_dev *pci_dev);
 void disable_ats_device(int seg, int bus, int devfn);
 struct pci_ats_dev *get_ats_device(int seg, int bus, int devfn);
 
diff --git a/xen/drivers/passthrough/vtd/iommu.c b/xen/drivers/passthrough/vtd/iommu.c
index f010612..1b0a0f0 100644
--- a/xen/drivers/passthrough/vtd/iommu.c
+++ b/xen/drivers/passthrough/vtd/iommu.c
@@ -1461,8 +1461,8 @@ int domain_context_mapping_one(
     return rc;
 }
 
-static int domain_context_mapping(
-    struct domain *domain, u8 devfn, const struct pci_dev *pdev)
+static int domain_context_mapping(struct domain *domain, u8 devfn,
+                                  struct pci_dev *pdev)
 {
     struct acpi_drhd_unit *drhd;
     int ret = 0;
@@ -1498,7 +1498,7 @@ static int domain_context_mapping(
         ret = domain_context_mapping_one(domain, drhd->iommu, bus, devfn,
                                          pdev);
         if ( !ret && devfn == pdev->devfn && ats_device(pdev, drhd) > 0 )
-            enable_ats_device(seg, bus, devfn, drhd->iommu);
+            enable_ats_device(drhd->iommu, pdev);
 
         break;
 
@@ -1994,7 +1994,7 @@ static int intel_iommu_enable_device(struct pci_dev *pdev)
     if ( ret <= 0 )
         return ret;
 
-    ret = enable_ats_device(pdev->seg, pdev->bus, pdev->devfn, drhd->iommu);
+    ret = enable_ats_device(drhd->iommu, pdev);
 
     return ret >= 0 ? 0 : ret;
 }
diff --git a/xen/drivers/passthrough/vtd/x86/ats.c b/xen/drivers/passthrough/vtd/x86/ats.c
index dfa4d30..a6c53ea 100644
--- a/xen/drivers/passthrough/vtd/x86/ats.c
+++ b/xen/drivers/passthrough/vtd/x86/ats.c
@@ -76,21 +76,25 @@ static int device_in_domain(struct iommu *iommu, struct pci_ats_dev *pdev, u16 d
     struct root_entry *root_entry = NULL;
     struct context_entry *ctxt_entry = NULL;
     int tt, found = 0;
+    struct pci_dev *pci_dev = pdev->pci_dev;
+
+    if ( !pci_dev )
+        return -ENODEV;
 
     root_entry = (struct root_entry *) map_vtd_domain_page(iommu->root_maddr);
-    if ( !root_entry || !root_present(root_entry[pdev->bus]) )
+    if ( !root_entry || !root_present(root_entry[pci_dev->bus]) )
         goto out;
 
     ctxt_entry = (struct context_entry *)
-                 map_vtd_domain_page(root_entry[pdev->bus].val);
+                 map_vtd_domain_page(root_entry[pci_dev->bus].val);
 
     if ( ctxt_entry == NULL )
         goto out;
 
-    if ( context_domain_id(ctxt_entry[pdev->devfn]) != did )
+    if ( context_domain_id(ctxt_entry[pci_dev->devfn]) != did )
         goto out;
 
-    tt = context_translation_type(ctxt_entry[pdev->devfn]);
+    tt = context_translation_type(ctxt_entry[pci_dev->devfn]);
     if ( tt != CONTEXT_TT_DEV_IOTLB )
         goto out;
 
@@ -116,10 +120,16 @@ int dev_invalidate_iotlb(struct iommu *iommu, u16 did,
 
     list_for_each_entry( pdev, &ats_devices, list )
     {
-        u16 sid = PCI_BDF2(pdev->bus, pdev->devfn);
+        struct pci_dev *pci_dev = pdev->pci_dev;
+        u16 sid;
         bool_t sbit;
         int rc = 0;
 
+        if ( !pci_dev )
+            continue;
+
+        sid = PCI_BDF2(pci_dev->bus, pci_dev->devfn);
+
         /* Only invalidate devices that belong to this IOMMU */
         if ( pdev->iommu != iommu )
             continue;
diff --git a/xen/drivers/passthrough/x86/ats.c b/xen/drivers/passthrough/x86/ats.c
index 40c9f40..8ce759f 100644
--- a/xen/drivers/passthrough/x86/ats.c
+++ b/xen/drivers/passthrough/x86/ats.c
@@ -22,26 +22,34 @@ LIST_HEAD(ats_devices);
 bool_t __read_mostly ats_enabled = 0;
 boolean_param("ats", ats_enabled);
 
-int enable_ats_device(int seg, int bus, int devfn, const void *iommu)
+int enable_ats_device(const void *iommu, struct pci_dev *pci_dev)
 {
     struct pci_ats_dev *pdev = NULL;
     u32 value;
     int pos;
 
-    pos = pci_find_ext_capability(seg, bus, devfn, PCI_EXT_CAP_ID_ATS);
+    pos = pci_find_ext_capability(pci_dev->seg, pci_dev->bus, pci_dev->devfn,
+                                  PCI_EXT_CAP_ID_ATS);
     BUG_ON(!pos);
 
     if ( iommu_verbose )
         dprintk(XENLOG_INFO, "%04x:%02x:%02x.%u: ATS capability found\n",
-                seg, bus, PCI_SLOT(devfn), PCI_FUNC(devfn));
+                pci_dev->seg, pci_dev->bus, PCI_SLOT(pci_dev->devfn),
+                PCI_FUNC(pci_dev->devfn));
 
-    value = pci_conf_read16(seg, bus, PCI_SLOT(devfn),
-                            PCI_FUNC(devfn), pos + ATS_REG_CTL);
+    value = pci_conf_read16(pci_dev->seg, pci_dev->bus, PCI_SLOT(pci_dev->devfn),
+                            PCI_FUNC(pci_dev->devfn), pos + ATS_REG_CTL);
     if ( value & ATS_ENABLE )
     {
         list_for_each_entry ( pdev, &ats_devices, list )
         {
-            if ( pdev->seg == seg && pdev->bus == bus && pdev->devfn == devfn )
+            struct pci_dev *pci_dev = pdev->pci_dev;
+
+            if ( !pci_dev )
+                continue;
+            if ( pci_dev->seg == pci_dev->seg &&
+                 pci_dev->bus == pci_dev->bus &&
+                 pci_dev->devfn == pci_dev->devfn )
             {
                 pos = 0;
                 break;
@@ -56,18 +64,16 @@ int enable_ats_device(int seg, int bus, int devfn, const void *iommu)
     if ( !(value & ATS_ENABLE) )
     {
         value |= ATS_ENABLE;
-        pci_conf_write16(seg, bus, PCI_SLOT(devfn), PCI_FUNC(devfn),
-                         pos + ATS_REG_CTL, value);
+        pci_conf_write16(pci_dev->seg, pci_dev->bus, PCI_SLOT(pci_dev->devfn),
+                         PCI_FUNC(pci_dev->devfn), pos + ATS_REG_CTL, value);
     }
 
     if ( pos )
     {
-        pdev->seg = seg;
-        pdev->bus = bus;
-        pdev->devfn = devfn;
+        pdev->pci_dev = pci_dev;
         pdev->iommu = iommu;
-        value = pci_conf_read16(seg, bus, PCI_SLOT(devfn),
-                                PCI_FUNC(devfn), pos + ATS_REG_CAP);
+        value = pci_conf_read16(pci_dev->seg, pci_dev->bus, PCI_SLOT(pci_dev->devfn),
+                                PCI_FUNC(pci_dev->devfn), pos + ATS_REG_CAP);
         pdev->ats_queue_depth = value & ATS_QUEUE_DEPTH_MASK ?:
                                 ATS_QUEUE_DEPTH_MASK + 1;
         list_add(&pdev->list, &ats_devices);
@@ -75,8 +81,8 @@ int enable_ats_device(int seg, int bus, int devfn, const void *iommu)
 
     if ( iommu_verbose )
         dprintk(XENLOG_INFO, "%04x:%02x:%02x.%u: ATS %s enabled\n",
-                seg, bus, PCI_SLOT(devfn), PCI_FUNC(devfn),
-                pos ? "is" : "was");
+                pci_dev->seg, pci_dev->bus, PCI_SLOT(pci_dev->devfn),
+                PCI_FUNC(pci_dev->devfn), pos ? "is" : "was");
 
     return pos;
 }
@@ -98,7 +104,13 @@ void disable_ats_device(int seg, int bus, int devfn)
 
     list_for_each_entry ( pdev, &ats_devices, list )
     {
-        if ( pdev->seg == seg && pdev->bus == bus && pdev->devfn == devfn )
+        struct pci_dev *pci_dev = pdev->pci_dev;
+
+        if ( !pci_dev )
+            continue;
+        if ( pci_dev->seg == seg &&
+             pci_dev->bus == bus &&
+             pci_dev->devfn == devfn )
         {
             list_del(&pdev->list);
             xfree(pdev);
@@ -120,7 +132,13 @@ struct pci_ats_dev *get_ats_device(int seg, int bus, int devfn)
 
     list_for_each_entry ( pdev, &ats_devices, list )
     {
-        if ( pdev->seg == seg && pdev->bus == bus && pdev->devfn == devfn )
+        struct pci_dev *pci_dev = pdev->pci_dev;
+
+        if ( !pci_dev )
+            continue;
+        if ( pci_dev->seg == seg &&
+             pci_dev->bus == bus &&
+             pci_dev->devfn == devfn )
             return pdev;
     }
 
