void (*drop_interrupt_shadow)(struct kvm_vcpu *vcpu);
        int (*set_tss_addr)(struct kvm *kvm, unsigned int addr);
        int (*get_tdp_level)(void);
-       int (*get_mt_mask_shift)(void);
+       u64 (*get_mt_mask)(struct kvm_vcpu *vcpu, gfn_t gfn, bool is_mmio);
 };
 
 extern struct kvm_x86_ops *kvm_x86_ops;
 void kvm_mmu_set_nonpresent_ptes(u64 trap_pte, u64 notrap_pte);
 void kvm_mmu_set_base_ptes(u64 base_pte);
 void kvm_mmu_set_mask_ptes(u64 user_mask, u64 accessed_mask,
-               u64 dirty_mask, u64 nx_mask, u64 x_mask, u64 mt_mask);
+               u64 dirty_mask, u64 nx_mask, u64 x_mask);
 
 int kvm_mmu_reset_context(struct kvm_vcpu *vcpu);
 void kvm_mmu_slot_remove_write_access(struct kvm *kvm, int slot);
                          const void *val, int bytes);
 int kvm_pv_mmu_op(struct kvm_vcpu *vcpu, unsigned long bytes,
                  gpa_t addr, unsigned long *ret);
+u8 kvm_get_guest_memory_type(struct kvm_vcpu *vcpu, gfn_t gfn);
 
 extern bool tdp_enabled;
 
 
 static u64 __read_mostly shadow_user_mask;
 static u64 __read_mostly shadow_accessed_mask;
 static u64 __read_mostly shadow_dirty_mask;
-static u64 __read_mostly shadow_mt_mask;
 
 static inline u64 rsvd_bits(int s, int e)
 {
 EXPORT_SYMBOL_GPL(kvm_mmu_set_base_ptes);
 
 void kvm_mmu_set_mask_ptes(u64 user_mask, u64 accessed_mask,
-               u64 dirty_mask, u64 nx_mask, u64 x_mask, u64 mt_mask)
+               u64 dirty_mask, u64 nx_mask, u64 x_mask)
 {
        shadow_user_mask = user_mask;
        shadow_accessed_mask = accessed_mask;
        shadow_dirty_mask = dirty_mask;
        shadow_nx_mask = nx_mask;
        shadow_x_mask = x_mask;
-       shadow_mt_mask = mt_mask;
 }
 EXPORT_SYMBOL_GPL(kvm_mmu_set_mask_ptes);
 
        return mtrr_state->def_type;
 }
 
-static u8 get_memory_type(struct kvm_vcpu *vcpu, gfn_t gfn)
+u8 kvm_get_guest_memory_type(struct kvm_vcpu *vcpu, gfn_t gfn)
 {
        u8 mtrr;
 
                mtrr = MTRR_TYPE_WRBACK;
        return mtrr;
 }
+EXPORT_SYMBOL_GPL(kvm_get_guest_memory_type);
 
 static int kvm_unsync_page(struct kvm_vcpu *vcpu, struct kvm_mmu_page *sp)
 {
 {
        u64 spte;
        int ret = 0;
-       u64 mt_mask = shadow_mt_mask;
 
        /*
         * We don't set the accessed bit, since we sometimes want to see
                spte |= shadow_user_mask;
        if (largepage)
                spte |= PT_PAGE_SIZE_MASK;
-       if (mt_mask) {
-               if (!kvm_is_mmio_pfn(pfn)) {
-                       mt_mask = get_memory_type(vcpu, gfn) <<
-                               kvm_x86_ops->get_mt_mask_shift();
-                       mt_mask |= VMX_EPT_IGMT_BIT;
-               } else
-                       mt_mask = MTRR_TYPE_UNCACHABLE <<
-                               kvm_x86_ops->get_mt_mask_shift();
-               spte |= mt_mask;
-       }
+       if (tdp_enabled)
+               spte |= kvm_x86_ops->get_mt_mask(vcpu, gfn,
+                       kvm_is_mmio_pfn(pfn));
 
        spte |= (u64)pfn << PAGE_SHIFT;
 
 
 #endif
 }
 
-static int svm_get_mt_mask_shift(void)
+static u64 svm_get_mt_mask(struct kvm_vcpu *vcpu, gfn_t gfn, bool is_mmio)
 {
        return 0;
 }
 
        .set_tss_addr = svm_set_tss_addr,
        .get_tdp_level = get_npt_level,
-       .get_mt_mask_shift = svm_get_mt_mask_shift,
+       .get_mt_mask = svm_get_mt_mask,
 };
 
 static int __init svm_init(void)
 
        return VMX_EPT_DEFAULT_GAW + 1;
 }
 
-static int vmx_get_mt_mask_shift(void)
+static u64 vmx_get_mt_mask(struct kvm_vcpu *vcpu, gfn_t gfn, bool is_mmio)
 {
-       return VMX_EPT_MT_EPTE_SHIFT;
+       u64 ret;
+
+       if (is_mmio)
+               ret = MTRR_TYPE_UNCACHABLE << VMX_EPT_MT_EPTE_SHIFT;
+       else
+               ret = (kvm_get_guest_memory_type(vcpu, gfn) <<
+                       VMX_EPT_MT_EPTE_SHIFT) | VMX_EPT_IGMT_BIT;
+
+       return ret;
 }
 
 static struct kvm_x86_ops vmx_x86_ops = {
 
        .set_tss_addr = vmx_set_tss_addr,
        .get_tdp_level = get_ept_level,
-       .get_mt_mask_shift = vmx_get_mt_mask_shift,
+       .get_mt_mask = vmx_get_mt_mask,
 };
 
 static int __init vmx_init(void)
                kvm_mmu_set_base_ptes(VMX_EPT_READABLE_MASK |
                        VMX_EPT_WRITABLE_MASK);
                kvm_mmu_set_mask_ptes(0ull, 0ull, 0ull, 0ull,
-                               VMX_EPT_EXECUTABLE_MASK,
-                               VMX_EPT_DEFAULT_MT << VMX_EPT_MT_EPTE_SHIFT);
+                               VMX_EPT_EXECUTABLE_MASK);
                kvm_enable_tdp();
        } else
                kvm_disable_tdp();
 
        kvm_mmu_set_nonpresent_ptes(0ull, 0ull);
        kvm_mmu_set_base_ptes(PT_PRESENT_MASK);
        kvm_mmu_set_mask_ptes(PT_USER_MASK, PT_ACCESSED_MASK,
-                       PT_DIRTY_MASK, PT64_NX_MASK, 0, 0);
+                       PT_DIRTY_MASK, PT64_NX_MASK, 0);
 
        for_each_possible_cpu(cpu)
                per_cpu(cpu_tsc_khz, cpu) = tsc_khz;