diff --git a/arch/x86/mm/fault_32.c b/arch/x86/mm/fault_32.c
index 276863dc4bdd45b88ee266b96f0046d98ea875af..93ede2dde958ed9ec82d7cc87b8c26f1a9f239cd 100644
--- a/arch/x86/mm/fault_32.c
+++ b/arch/x86/mm/fault_32.c
@@ -475,14 +475,16 @@ void __kprobes do_page_fault(struct pt_regs *regs, unsigned long error_code)
 	else
 		tsk->min_flt++;
 
+#ifdef CONFIG_X86_32
 	/*
 	 * Did it hit the DOS screen memory VA from vm86 mode?
 	 */
-	if (regs->flags & VM_MASK) {
+	if (v8086_mode(regs)) {
 		unsigned long bit = (address - 0xA0000) >> PAGE_SHIFT;
 		if (bit < 32)
 			tsk->thread.screen_bitmap |= 1 << bit;
 	}
+#endif
 	up_read(&mm->mmap_sem);
 	return;
 
diff --git a/arch/x86/mm/fault_64.c b/arch/x86/mm/fault_64.c
index 9ef0306efe9ef0023383525da351805b8dc575d0..2d9e6da2cb30622ea4f7b6915d72b99ff709e344 100644
--- a/arch/x86/mm/fault_64.c
+++ b/arch/x86/mm/fault_64.c
@@ -514,6 +514,17 @@ asmlinkage void __kprobes do_page_fault(struct pt_regs *regs,
 		tsk->maj_flt++;
 	else
 		tsk->min_flt++;
+
+#ifdef CONFIG_X86_32
+	/*
+	 * Did it hit the DOS screen memory VA from vm86 mode?
+	 */
+	if (v8086_mode(regs)) {
+		unsigned long bit = (address - 0xA0000) >> PAGE_SHIFT;
+		if (bit < 32)
+			tsk->thread.screen_bitmap |= 1 << bit;
+	}
+#endif
 	up_read(&mm->mmap_sem);
 	return;