diff --git a/src/arch/x86_64/cls.c b/src/arch/x86_64/cls.c
index 802b9d1c..973dd496 100644
--- a/src/arch/x86_64/cls.c
+++ b/src/arch/x86_64/cls.c
@@ -5,12 +5,12 @@
 cls_t bspcls;
 
 
-void arch_setcls(void* addr){
+void arch_setcls(cls_t* addr){
 	wrmsr(MSR_GSBASE, (uint64_t)addr);
 }
 
-void* arch_getcls(){
-	return (void*)rdmsr(MSR_GSBASE);
+cls_t* arch_getcls(){
+	return (cls_t*)rdmsr(MSR_GSBASE);
 }
 
 void bsp_setcls(){
diff --git a/src/arch/x86_64/include/arch/cls.h b/src/arch/x86_64/include/arch/cls.h
index 09433e71..a4266997 100644
--- a/src/arch/x86_64/include/arch/cls.h
+++ b/src/arch/x86_64/include/arch/cls.h
@@ -5,7 +5,7 @@
 #include <arch/ist.h>
 #include <arch/ist.h>
 #include <arch/regs.h>
-#include <arch/mmu.h>
+#include <kernel/vmm.h>
 
 // cpu level storage
 // this will be pointed to by GS and will contain per cpu info
@@ -14,11 +14,12 @@ typedef struct{
 	gdt_t gdt;
 	ist_t ist;
 	arch_regserror *laststate;
+	vmm_context *context;
 } cls_t;
 
 void bsp_setcls();
-void arch_setcls(void* addr);
-void* arch_getcls();
+void arch_setcls(cls_t* addr);
+cls_t* arch_getcls();
 
 static cls_t __seg_gs *cls = 0;