diff --git a/inc/hf/vm.h b/inc/hf/vm.h
index dc72e4a..6bc4882 100644
--- a/inc/hf/vm.h
+++ b/inc/hf/vm.h
@@ -52,4 +52,3 @@
 uint32_t vm_get_count(void);
 struct vm *vm_get(uint32_t id);
 void vm_start_vcpu(struct vm *vm, size_t index, ipaddr_t entry, uintreg_t arg);
-void vm_set_current(struct vm *vm);
diff --git a/src/api.c b/src/api.c
index 8761e61..16df0d9 100644
--- a/src/api.c
+++ b/src/api.c
@@ -41,9 +41,6 @@
 	struct vm *primary = vm_get(HF_PRIMARY_VM_ID);
 	struct vcpu *next = &primary->vcpus[cpu_index(current->cpu)];
 
-	/* Switch back to primary VM. */
-	vm_set_current(primary);
-
 	/* Set the return value for the primary VM's call to HF_VCPU_RUN. */
 	arch_regs_set_retval(&next->regs,
 			     hf_vcpu_run_return_encode(primary_ret));
@@ -151,7 +148,6 @@
 	} else {
 		vcpu->cpu = current->cpu;
 		vcpu->state = vcpu_state_running;
-		vm_set_current(vm);
 		*next = vcpu;
 		ret.code = HF_VCPU_RUN_YIELD;
 	}
diff --git a/src/arch/aarch64/exceptions.S b/src/arch/aarch64/exceptions.S
index 4e3219f..de349a1 100644
--- a/src/arch/aarch64/exceptions.S
+++ b/src/arch/aarch64/exceptions.S
@@ -279,7 +279,15 @@
 	stp x20, x21, [x1, #16 * 10]
 
 	mrs x22, par_el1
-	str x22, [x1, #16 * 11]
+	mrs x23, hcr_el2
+	stp x22, x23, [x1, #16 * 11]
+
+	mrs x24, cptr_el2
+	mrs x25, cnthctl_el2
+	stp x24, x25, [x1, #16 * 12]
+
+	mrs x26, vttbr_el2
+	str x26, [x1, #16 * 13]
 
 	/* Intentional fallthrough. */
 
@@ -335,8 +343,16 @@
 	msr sp_el0, x20
 	msr sp_el1, x21
 
-	ldr x22, [x0, #16 * 11]
+	ldp x22, x23, [x0, #16 * 11]
 	msr par_el1, x22
+	msr hcr_el2, x23
+
+	ldp x24, x25, [x0, #16 * 12]
+	msr cptr_el2, x22
+	msr cnthctl_el2, x23
+
+	ldr x26, [x0, #16 * 13]
+	msr vttbr_el2, x26
 
 	/* Restore non-volatile registers. */
 	add x0, x0, #(VCPU_REGS - VCPU_LAZY)
diff --git a/src/arch/aarch64/inc/hf/arch/cpu.h b/src/arch/aarch64/inc/hf/arch/cpu.h
index 78f3cf2..efbad6e 100644
--- a/src/arch/aarch64/inc/hf/arch/cpu.h
+++ b/src/arch/aarch64/inc/hf/arch/cpu.h
@@ -55,6 +55,10 @@
 		uintreg_t sp_el0;
 		uintreg_t sp_el1;
 		uintreg_t par_el1;
+		uintreg_t hcr_el2;
+		uintreg_t cptr_el2;
+		uintreg_t cnthctl_el2;
+		uintreg_t vttbr_el2;
 	} lazy;
 };
 
@@ -68,7 +72,9 @@
 	__asm__ volatile("msr DAIFClr, #0xf");
 }
 
-static inline void arch_cpu_update(bool is_primary)
+static inline void arch_regs_init(struct arch_regs *r, bool is_primary,
+				  uint64_t vmid, paddr_t table, ipaddr_t pc,
+				  uintreg_t arg)
 {
 	uintreg_t hcr;
 	uintreg_t cptr;
@@ -97,14 +103,10 @@
 			   (1u << 1);  /* EL1PCEN, trap phys timer access. */
 	}
 
-	__asm__ volatile("msr hcr_el2, %0" ::"r"(hcr));
-	__asm__ volatile("msr cptr_el2, %0" ::"r"(cptr));
-	__asm__ volatile("msr cnthctl_el2, %0" ::"r"(cnthctl));
-}
-
-static inline void arch_regs_init(struct arch_regs *r, ipaddr_t pc,
-				  uintreg_t arg)
-{
+	r->lazy.hcr_el2 = hcr;
+	r->lazy.cptr_el2 = cptr;
+	r->lazy.cnthctl_el2 = cnthctl;
+	r->lazy.vttbr_el2 = pa_addr(table) | (vmid << 48);
 	/* TODO: Use constant here. */
 	r->spsr = 5 |	 /* M bits, set to EL1h. */
 		  (0xf << 6); /* DAIF bits set; disable interrupts. */
diff --git a/src/arch/aarch64/inc/hf/arch/mm.h b/src/arch/aarch64/inc/hf/arch/mm.h
index 7ee6d68..35278ca 100644
--- a/src/arch/aarch64/inc/hf/arch/mm.h
+++ b/src/arch/aarch64/inc/hf/arch/mm.h
@@ -193,13 +193,6 @@
 		"dsb ish\n");
 }
 
-static inline void arch_mm_set_vm(uint64_t vmid, paddr_t table)
-{
-	__asm__ volatile("msr vttbr_el2, %0"
-			 :
-			 : "r"(pa_addr(table) | (vmid << 48)));
-}
-
 uint64_t arch_mm_mode_to_attrs(int mode);
 bool arch_mm_init(paddr_t table, bool first);
 uint8_t arch_mm_max_level(int mode);
diff --git a/src/arch/fake/inc/hf/arch/cpu.h b/src/arch/fake/inc/hf/arch/cpu.h
index d374bcb..d8ede27 100644
--- a/src/arch/fake/inc/hf/arch/cpu.h
+++ b/src/arch/fake/inc/hf/arch/cpu.h
@@ -37,16 +37,14 @@
 	/* TODO */
 }
 
-static inline void arch_cpu_update(bool is_primary)
-{
-	/* TODO */
-	(void)is_primary;
-}
-
-static inline void arch_regs_init(struct arch_regs *r, ipaddr_t pc,
+static inline void arch_regs_init(struct arch_regs *r, bool is_primary,
+				  uint64_t vmid, paddr_t table, ipaddr_t pc,
 				  uintreg_t arg)
 {
 	/* TODO */
+	(void)is_primary;
+	(void)vmid;
+	(void)table;
 	(void)pc;
 	r->r[0] = arg;
 }
diff --git a/src/arch/fake/inc/hf/arch/mm.h b/src/arch/fake/inc/hf/arch/mm.h
index 98c73c7..c99bf92 100644
--- a/src/arch/fake/inc/hf/arch/mm.h
+++ b/src/arch/fake/inc/hf/arch/mm.h
@@ -203,10 +203,3 @@
 	(void)first;
 	return true;
 }
-
-static inline void arch_mm_set_vm(uint64_t vmid, paddr_t table)
-{
-	/* TODO */
-	(void)vmid;
-	(void)table;
-}
diff --git a/src/cpu.c b/src/cpu.c
index 9c7484b..4c6a18b 100644
--- a/src/cpu.c
+++ b/src/cpu.c
@@ -92,9 +92,10 @@
 	sl_unlock(&c->lock);
 
 	if (!prev) {
-		struct vcpu *vcpu =
-			&vm_get(HF_PRIMARY_VM_ID)->vcpus[cpu_index(c)];
-		arch_regs_init(&vcpu->regs, entry, arg);
+		struct vm *vm = vm_get(HF_PRIMARY_VM_ID);
+		struct vcpu *vcpu = &vm->vcpus[cpu_index(c)];
+		arch_regs_init(&vcpu->regs, true, vm->id, vm->ptable.table,
+			       entry, arg);
 		vcpu_on(vcpu);
 	}
 
diff --git a/src/main.c b/src/main.c
index db191a8..eb233d9 100644
--- a/src/main.c
+++ b/src/main.c
@@ -130,7 +130,6 @@
  */
 struct vcpu *cpu_main(struct cpu *c)
 {
-	struct vm *primary;
 	struct vcpu *vcpu;
 
 	/*
@@ -149,10 +148,7 @@
 		panic("mm_cpu_init failed");
 	}
 
-	primary = vm_get(HF_PRIMARY_VM_ID);
-	vm_set_current(primary);
-
-	vcpu = &primary->vcpus[cpu_index(c)];
+	vcpu = &vm_get(HF_PRIMARY_VM_ID)->vcpus[cpu_index(c)];
 	vcpu->cpu = c;
 	return vcpu;
 }
diff --git a/src/vm.c b/src/vm.c
index 02b60ff..7387f20 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -73,13 +73,8 @@
 {
 	struct vcpu *vcpu = &vm->vcpus[index];
 	if (index < vm->vcpu_count) {
-		arch_regs_init(&vcpu->regs, entry, arg);
+		arch_regs_init(&vcpu->regs, vm->id == HF_PRIMARY_VM_ID, vm->id,
+			       vm->ptable.table, entry, arg);
 		vcpu_on(vcpu);
 	}
 }
-
-void vm_set_current(struct vm *vm)
-{
-	arch_cpu_update(vm->id == HF_PRIMARY_VM_ID);
-	arch_mm_set_vm(vm->id, vm->ptable.table);
-}