diff --git a/Makefile b/Makefile
index b8a2071b..4afa1b88 100644
--- a/Makefile
+++ b/Makefile
@@ -124,6 +124,10 @@ ifeq ($(CONFIG_ACPICA),y)
 COMMON_FLAGS += -DKTF_ACPICA
 endif
 
+ifneq ($(UNITTEST),)
+COMMON_FLAGS += -DKTF_UNIT_TEST
+endif
+
 AFLAGS  := $(COMMON_FLAGS) -DASM_FILE -D__ASSEMBLY__ -nostdlib -nostdinc
 CFLAGS  := $(COMMON_FLAGS) -std=gnu99 -O2 -g -Wall -Wextra -ffreestanding -nostdlib -nostdinc
 CFLAGS  += -mno-red-zone -mno-mmx -mno-sse -mno-sse2
diff --git a/arch/x86/real_mode.S b/arch/x86/real_mode.S
index c42c11de..a2843901 100644
--- a/arch/x86/real_mode.S
+++ b/arch/x86/real_mode.S
@@ -139,7 +139,18 @@ END_FUNC(_prot_to_real)
 
 .align 16
 .Lfrom_long_mode:
-    /* Disable LME in EFER */
+    /* Use protected mode data segment */
+    mov   $__KERN_DS32, %eax
+    mov   %eax, %ds
+    mov   %eax, %ss
+
+    /* Disable paging to enter protected mode. */
+    /* Clearing PG implicitly clears EFER.LMA, see Intel SDM 9.8.5.4, */
+    mov   %cr0, %eax
+    and   $~(X86_CR0_PG | X86_CR0_WP), %eax
+    mov   %eax, %cr0
+
+    /* Disable LME in EFER. */
     movl  $MSR_EFER, %ecx
     rdmsr
     and   $~EFER_LME, %eax
@@ -151,15 +162,6 @@ END_FUNC(_prot_to_real)
     lgdt  rmode_gdt_ptr
     lidt  rmode_idt_ptr
 
-    /* Disable paging to enter protected mode */
-    mov   %cr0, %eax
-    and   $~(X86_CR0_PG | X86_CR0_WP), %eax
-    mov   %eax, %cr0
-
-    /* Use protected mode data selector */
-    mov   $__KERN_DS32, %eax
-    mov   %eax, %ds
-
     /* Use real mode accessible stack */
     mov   rmode_stack, %esp
 
diff --git a/include/arch/x86/asm-macros.h b/include/arch/x86/asm-macros.h
index 17a2e415..a4159e54 100644
--- a/include/arch/x86/asm-macros.h
+++ b/include/arch/x86/asm-macros.h
@@ -69,6 +69,26 @@
 .endm
 
 .macro SAVE_ALL_REGS32
+    push %eax
+    push %ebx
+    push %ecx
+    push %edx
+    push %esi
+    push %edi
+    push %ebp
+.endm
+
+.macro RESTORE_ALL_REGS32
+    pop %ebp
+    pop %edi
+    pop %esi
+    pop %edx
+    pop %ecx
+    pop %ebx
+    pop %eax
+.endm
+
+.macro SAVE_ALL_REGS
     push %_ASM_AX
     push %_ASM_BX
     push %_ASM_CX
@@ -76,20 +96,6 @@
     push %_ASM_SI
     push %_ASM_DI
     push %_ASM_BP
-.endm
-
-.macro RESTORE_ALL_REGS32
-    pop %_ASM_BP
-    pop %_ASM_DI
-    pop %_ASM_SI
-    pop %_ASM_DX
-    pop %_ASM_CX
-    pop %_ASM_BX
-    pop %_ASM_AX
-.endm
-
-.macro SAVE_ALL_REGS
-    SAVE_ALL_REGS32
 #if defined(__x86_64__)
     push %r8
     push %r9
@@ -113,7 +119,13 @@
     pop %r9
     pop %r8
 #endif
-    RESTORE_ALL_REGS32
+    pop %_ASM_BP
+    pop %_ASM_DI
+    pop %_ASM_SI
+    pop %_ASM_DX
+    pop %_ASM_CX
+    pop %_ASM_BX
+    pop %_ASM_AX
 .endm
 
 .macro SAVE_CALLEE_SAVED_REGS
diff --git a/tests/unittests.c b/tests/unittests.c
index 954511f1..7fa4b014 100644
--- a/tests/unittests.c
+++ b/tests/unittests.c
@@ -173,10 +173,8 @@ int unit_tests(void *_unused) {
     schedule_task(task_user1, get_bsp_cpu());
     schedule_task(task_user2, get_cpu(1));
 
-#ifdef UNITTEST_LONGMODE
     printk("Long mode to real mode transition:\n");
     long_to_real();
-#endif
 
     return 0;
 }