return re->hi & VTD_PAGE_MASK;
 }
 
-static inline void context_clear_pasid_enable(struct context_entry *context)
-{
-       context->lo &= ~(1ULL << 11);
-}
-
-static inline bool context_pasid_enabled(struct context_entry *context)
-{
-       return !!(context->lo & (1ULL << 11));
-}
-
-static inline void context_set_copied(struct context_entry *context)
-{
-       context->hi |= (1ull << 3);
-}
-
-static inline bool context_copied(struct context_entry *context)
-{
-       return !!(context->hi & (1ULL << 3));
-}
-
-static inline bool __context_present(struct context_entry *context)
-{
-       return (context->lo & 1);
-}
-
-bool context_present(struct context_entry *context)
-{
-       return context_pasid_enabled(context) ?
-            __context_present(context) :
-            __context_present(context) && !context_copied(context);
-}
-
 static inline void context_set_present(struct context_entry *context)
 {
        context->lo |= 1;
        context->hi = 0;
 }
 
+static inline bool context_copied(struct intel_iommu *iommu, u8 bus, u8 devfn)
+{
+       if (!iommu->copied_tables)
+               return false;
+
+       return test_bit(((long)bus << 8) | devfn, iommu->copied_tables);
+}
+
+static inline void
+set_context_copied(struct intel_iommu *iommu, u8 bus, u8 devfn)
+{
+       set_bit(((long)bus << 8) | devfn, iommu->copied_tables);
+}
+
+static inline void
+clear_context_copied(struct intel_iommu *iommu, u8 bus, u8 devfn)
+{
+       clear_bit(((long)bus << 8) | devfn, iommu->copied_tables);
+}
+
 /*
  * This domain is a statically identity mapping domain.
  *     1. This domain creats a static 1:1 mapping to all usable memory.
        struct context_entry *context;
        u64 *entry;
 
+       /*
+        * Except that the caller requested to allocate a new entry,
+        * returning a copied context entry makes no sense.
+        */
+       if (!alloc && context_copied(iommu, bus, devfn))
+               return NULL;
+
        entry = &root->lo;
        if (sm_supported(iommu)) {
                if (devfn >= 0x80) {
                iommu->domain_ids = NULL;
        }
 
+       if (iommu->copied_tables) {
+               bitmap_free(iommu->copied_tables);
+               iommu->copied_tables = NULL;
+       }
+
        g_iommus[iommu->seq_id] = NULL;
 
        /* free context mapping */
                goto out_unlock;
 
        ret = 0;
-       if (context_present(context))
+       if (context_present(context) && !context_copied(iommu, bus, devfn))
                goto out_unlock;
 
        /*
         * in-flight DMA will exist, and we don't need to worry anymore
         * hereafter.
         */
-       if (context_copied(context)) {
+       if (context_copied(iommu, bus, devfn)) {
                u16 did_old = context_domain_id(context);
 
                if (did_old < cap_ndoms(iommu->cap)) {
                        iommu->flush.flush_iotlb(iommu, did_old, 0, 0,
                                                 DMA_TLB_DSI_FLUSH);
                }
+
+               clear_context_copied(iommu, bus, devfn);
        }
 
        context_clear_entry(context);
                /* Now copy the context entry */
                memcpy(&ce, old_ce + idx, sizeof(ce));
 
-               if (!__context_present(&ce))
+               if (!context_present(&ce))
                        continue;
 
                did = context_domain_id(&ce);
                if (did >= 0 && did < cap_ndoms(iommu->cap))
                        set_bit(did, iommu->domain_ids);
 
-               /*
-                * We need a marker for copied context entries. This
-                * marker needs to work for the old format as well as
-                * for extended context entries.
-                *
-                * Bit 67 of the context entry is used. In the old
-                * format this bit is available to software, in the
-                * extended format it is the PGE bit, but PGE is ignored
-                * by HW if PASIDs are disabled (and thus still
-                * available).
-                *
-                * So disable PASIDs first and then mark the entry
-                * copied. This means that we don't copy PASID
-                * translations from the old kernel, but this is fine as
-                * faults there are not fatal.
-                */
-               context_clear_pasid_enable(&ce);
-               context_set_copied(&ce);
-
+               set_context_copied(iommu, bus, devfn);
                new_ce[idx] = ce;
        }
 
        bool new_ext, ext;
 
        rtaddr_reg = dmar_readq(iommu->reg + DMAR_RTADDR_REG);
-       ext        = !!(rtaddr_reg & DMA_RTADDR_RTT);
-       new_ext    = !!ecap_ecs(iommu->ecap);
+       ext        = !!(rtaddr_reg & DMA_RTADDR_SMT);
+       new_ext    = !!sm_supported(iommu);
 
        /*
         * The RTT bit can only be changed when translation is disabled,
        if (new_ext != ext)
                return -EINVAL;
 
+       iommu->copied_tables = bitmap_zalloc(BIT_ULL(16), GFP_KERNEL);
+       if (!iommu->copied_tables)
+               return -ENOMEM;
+
        old_rt_phys = rtaddr_reg & VTD_PAGE_MASK;
        if (!old_rt_phys)
                return -EINVAL;
 
 #define ecap_dis(e)            (((e) >> 27) & 0x1)
 #define ecap_nest(e)           (((e) >> 26) & 0x1)
 #define ecap_mts(e)            (((e) >> 25) & 0x1)
-#define ecap_ecs(e)            (((e) >> 24) & 0x1)
 #define ecap_iotlb_offset(e)   ((((e) >> 8) & 0x3ff) * 16)
 #define ecap_max_iotlb_offset(e) (ecap_iotlb_offset(e) + 16)
 #define ecap_coherent(e)       ((e) & 0x1)
 #define DMA_GSTS_CFIS (((u32)1) << 23)
 
 /* DMA_RTADDR_REG */
-#define DMA_RTADDR_RTT (((u64)1) << 11)
 #define DMA_RTADDR_SMT (((u64)1) << 10)
 
 /* CCMD_REG */
 #ifdef CONFIG_INTEL_IOMMU
        unsigned long   *domain_ids; /* bitmap of domains */
        struct dmar_domain ***domains; /* ptr to domains */
+       unsigned long   *copied_tables; /* bitmap of copied tables */
        spinlock_t      lock; /* protect context, domain ids */
        struct root_entry *root_entry; /* virtual address */
 
        return !((unsigned long)pte & ~VTD_PAGE_MASK);
 }
 
+static inline bool context_present(struct context_entry *context)
+{
+       return (context->lo & 1);
+}
+
 extern struct dmar_drhd_unit * dmar_find_matched_drhd_unit(struct pci_dev *dev);
 extern int dmar_find_matched_atsr_unit(struct pci_dev *dev);
 
 #endif /* CONFIG_INTEL_IOMMU_DEBUGFS */
 
 extern const struct attribute_group *intel_iommu_groups[];
-bool context_present(struct context_entry *context);
 struct context_entry *iommu_context_addr(struct intel_iommu *iommu, u8 bus,
                                         u8 devfn, int alloc);