--- 2011-09-20.orig/xen/drivers/passthrough/vtd/dmar.c	2011-09-21 17:55:19.000000000 +0200
+++ 2011-09-20/xen/drivers/passthrough/vtd/dmar.c	2011-10-11 13:35:05.000000000 +0200
@@ -160,7 +160,7 @@ static int __init acpi_register_atsr_uni
     return 0;
 }
 
-struct acpi_drhd_unit * acpi_find_matched_drhd_unit(struct pci_dev *pdev)
+struct acpi_drhd_unit *acpi_find_matched_drhd_unit(const struct pci_dev *pdev)
 {
     u8 bus, devfn;
     struct acpi_drhd_unit *drhd;
@@ -204,19 +204,26 @@ struct acpi_drhd_unit * acpi_find_matche
     return include_all;
 }
 
-struct acpi_atsr_unit * acpi_find_matched_atsr_unit(u16 seg, u8 bus, u8 devfn)
+struct acpi_atsr_unit *acpi_find_matched_atsr_unit(const struct pci_dev *pdev)
 {
     struct acpi_atsr_unit *atsr;
     struct acpi_atsr_unit *all_ports = NULL;
+    u16 bdf = PCI_BDF2(pdev->bus, pdev->devfn);
 
     list_for_each_entry ( atsr, &acpi_atsr_units, list )
     {
-        if ( atsr->segment != seg )
+        unsigned int i;
+
+        if ( atsr->segment != pdev->seg )
             continue;
 
-        if ( test_bit(bus, atsr->scope.buses) )
+        if ( test_bit(pdev->bus, atsr->scope.buses) )
             return atsr;
 
+        for ( i = 0; i < atsr->scope.devices_cnt; ++i )
+            if ( atsr->scope.devices[i] == bdf )
+                return atsr;
+
         if ( atsr->all_ports )
             all_ports = atsr;
     }
--- 2011-09-20.orig/xen/drivers/passthrough/vtd/dmar.h	2011-08-25 15:06:43.000000000 +0200
+++ 2011-09-20/xen/drivers/passthrough/vtd/dmar.h	2011-10-11 13:30:30.000000000 +0200
@@ -86,8 +86,8 @@ struct acpi_rhsa_unit {
         for (idx = 0; (bdf = rmrr->scope.devices[idx]) && \
                  idx < rmrr->scope.devices_cnt; idx++)
 
-struct acpi_drhd_unit * acpi_find_matched_drhd_unit(struct pci_dev *pdev);
-struct acpi_atsr_unit * acpi_find_matched_atsr_unit(u16 seg, u8 bus, u8 devfn);
+struct acpi_drhd_unit *acpi_find_matched_drhd_unit(const struct pci_dev *);
+struct acpi_atsr_unit *acpi_find_matched_atsr_unit(const struct pci_dev *);
 
 #define DMAR_TYPE 1
 #define RMRR_TYPE 2
--- 2011-09-20.orig/xen/drivers/passthrough/vtd/extern.h	2011-04-11 08:33:59.000000000 +0200
+++ 2011-09-20/xen/drivers/passthrough/vtd/extern.h	2011-10-11 13:56:00.000000000 +0200
@@ -61,10 +61,9 @@ extern bool_t ats_enabled;
 
 struct acpi_drhd_unit * find_ats_dev_drhd(struct iommu *iommu);
 
-int ats_device(int seg, int bus, int devfn);
+int ats_device(const struct pci_dev *, const struct acpi_drhd_unit *);
 int enable_ats_device(int seg, int bus, int devfn);
 void disable_ats_device(int seg, int bus, int devfn);
-int invalidate_ats_tcs(struct iommu *iommu);
 
 int dev_invalidate_iotlb(struct iommu *iommu, u16 did,
                          u64 addr, unsigned int size_order, u64 type);
@@ -76,7 +75,8 @@ static inline struct acpi_drhd_unit *fin
     return NULL;
 }
 
-static inline int ats_device(int seg, int bus, int devfn)
+static inline int ats_device(const struct pci_dev *pdev,
+                             const struct acpi_drhd_unit *drhd)
 {
     return 0;
 }
--- 2011-09-20.orig/xen/drivers/passthrough/vtd/iommu.c	2011-09-21 17:56:24.000000000 +0200
+++ 2011-09-20/xen/drivers/passthrough/vtd/iommu.c	2011-10-11 14:03:23.000000000 +0200
@@ -1411,7 +1411,7 @@ static int domain_context_mapping(
                     domain->domain_id, seg, bus,
                     PCI_SLOT(devfn), PCI_FUNC(devfn));
         ret = domain_context_mapping_one(domain, drhd->iommu, bus, devfn);
-        if ( !ret && ats_device(seg, bus, devfn) )
+        if ( !ret && ats_device(pdev, drhd) > 0 )
             enable_ats_device(seg, bus, devfn);
 
         break;
@@ -1542,7 +1542,7 @@ static int domain_context_unmap(
                     domain->domain_id, seg, bus,
                     PCI_SLOT(devfn), PCI_FUNC(devfn));
         ret = domain_context_unmap_one(domain, iommu, bus, devfn);
-        if ( !ret && ats_device(seg, bus, devfn) )
+        if ( !ret && ats_device(pdev, drhd) > 0 )
             disable_ats_device(seg, bus, devfn);
 
         break;
--- 2011-09-20.orig/xen/drivers/passthrough/vtd/x86/ats.c	2011-08-25 15:13:05.000000000 +0200
+++ 2011-09-20/xen/drivers/passthrough/vtd/x86/ats.c	2011-10-11 14:03:00.000000000 +0200
@@ -83,40 +83,32 @@ struct acpi_drhd_unit * find_ats_dev_drh
     return NULL;
 }
 
-int ats_device(int seg, int bus, int devfn)
+int ats_device(const struct pci_dev *pdev, const struct acpi_drhd_unit *drhd)
 {
-    struct acpi_drhd_unit *drhd, *ats_drhd, *new_drhd;
-    struct pci_dev *pdev;
-    int pos = 0;
+    struct acpi_drhd_unit *ats_drhd;
+    int pos;
 
     if ( !ats_enabled || !iommu_qinval )
         return 0;
 
-    pdev = pci_get_pdev(seg, bus, devfn);
-    if ( !pdev )
-        return 0;
-
-    drhd = acpi_find_matched_drhd_unit(pdev);
-    if ( !drhd )
-        return 0;
-
     if ( !ecap_queued_inval(drhd->iommu->ecap) ||
          !ecap_dev_iotlb(drhd->iommu->ecap) )
         return 0;
 
-    if ( !acpi_find_matched_atsr_unit(seg, bus, devfn) )
+    if ( !acpi_find_matched_atsr_unit(pdev) )
         return 0;
 
     ats_drhd = find_ats_dev_drhd(drhd->iommu);
-    pos = pci_find_ext_capability(seg, bus, devfn, PCI_EXT_CAP_ID_ATS);
+    pos = pci_find_ext_capability(pdev->seg, pdev->bus, pdev->devfn,
+                                  PCI_EXT_CAP_ID_ATS);
 
     if ( pos && (ats_drhd == NULL) )
     {
-        new_drhd = xmalloc(struct acpi_drhd_unit);
-        if ( !new_drhd )
-            return 0;
-        memcpy(new_drhd, drhd, sizeof(struct acpi_drhd_unit));
-        list_add_tail(&new_drhd->list, &ats_dev_drhd_units);
+        ats_drhd = xmalloc(struct acpi_drhd_unit);
+        if ( !ats_drhd )
+            return -ENOMEM;
+        *ats_drhd = *drhd;
+        list_add_tail(&ats_drhd->list, &ats_dev_drhd_units);
     }
     return pos;
 }