diff --git a/driver/linux b/driver/linux
index 01390ae..ddc3394 160000
--- a/driver/linux
+++ b/driver/linux
@@ -1 +1 @@
-Subproject commit 01390ae4e344f1363f1c169e3430fa3e72cafaee
+Subproject commit ddc3394d6462cd4e4a33ea3dc7edca7aed0f1ca3
diff --git a/inc/hf/api.h b/inc/hf/api.h
index 7d6b7c0..4060075 100644
--- a/inc/hf/api.h
+++ b/inc/hf/api.h
@@ -27,7 +27,8 @@
 spci_vm_count_t api_vm_get_count(void);
 int64_t api_vcpu_get_count(spci_vm_id_t vm_id, const struct vcpu *current);
 void api_regs_state_saved(struct vcpu *vcpu);
-struct hf_vcpu_run_return api_vcpu_run(spci_vm_id_t vm_id, uint32_t vcpu_idx,
+struct hf_vcpu_run_return api_vcpu_run(spci_vm_id_t vm_id,
+				       spci_vcpu_index_t vcpu_idx,
 				       const struct vcpu *current,
 				       struct vcpu **next);
 int64_t api_vm_configure(ipaddr_t send, ipaddr_t recv, struct vcpu *current,
@@ -47,7 +48,7 @@
 int64_t api_interrupt_enable(uint32_t intid, bool enable, struct vcpu *current);
 uint32_t api_interrupt_get(struct vcpu *current);
 int64_t api_interrupt_inject(spci_vm_id_t target_vm_id,
-			     uint32_t target_vcpu_idx, uint32_t intid,
+			     spci_vcpu_index_t target_vcpu_idx, uint32_t intid,
 			     struct vcpu *current, struct vcpu **next);
 
 int32_t api_spci_msg_send(uint32_t attributes, struct vcpu *current,
diff --git a/inc/hf/cpu.h b/inc/hf/cpu.h
index d0a5466..d370830 100644
--- a/inc/hf/cpu.h
+++ b/inc/hf/cpu.h
@@ -132,7 +132,7 @@
 void vcpu_unlock(struct vcpu_locked *locked);
 void vcpu_init(struct vcpu *vcpu, struct vm *vm);
 void vcpu_on(struct vcpu_locked vcpu, ipaddr_t entry, uintreg_t arg);
-size_t vcpu_index(const struct vcpu *vcpu);
+spci_vcpu_index_t vcpu_index(const struct vcpu *vcpu);
 bool vcpu_is_off(struct vcpu_locked vcpu);
 bool vcpu_secondary_reset_and_start(struct vcpu *vcpu, ipaddr_t entry,
 				    uintreg_t arg);
diff --git a/inc/hf/vm.h b/inc/hf/vm.h
index 4bebc20..bd8a905 100644
--- a/inc/hf/vm.h
+++ b/inc/hf/vm.h
@@ -102,4 +102,4 @@
 struct vm *vm_find(spci_vm_id_t id);
 struct vm_locked vm_lock(struct vm *vm);
 void vm_unlock(struct vm_locked *locked);
-struct vcpu *vm_get_vcpu(struct vm *vm, uint32_t vcpu_index);
+struct vcpu *vm_get_vcpu(struct vm *vm, spci_vcpu_index_t vcpu_index);
diff --git a/inc/vmapi/hf/abi.h b/inc/vmapi/hf/abi.h
index a625a8e..538a42e 100644
--- a/inc/vmapi/hf/abi.h
+++ b/inc/vmapi/hf/abi.h
@@ -90,7 +90,7 @@
 	union {
 		struct {
 			spci_vm_id_t vm_id;
-			uint16_t vcpu;
+			spci_vcpu_index_t vcpu;
 		} wake_up;
 		struct {
 			spci_vm_id_t vm_id;
diff --git a/inc/vmapi/hf/call.h b/inc/vmapi/hf/call.h
index 01be71f..64e658f 100644
--- a/inc/vmapi/hf/call.h
+++ b/inc/vmapi/hf/call.h
@@ -75,7 +75,7 @@
  * Returns an hf_vcpu_run_return struct telling the scheduler what to do next.
  */
 static inline struct hf_vcpu_run_return hf_vcpu_run(spci_vm_id_t vm_id,
-						    uint32_t vcpu_idx)
+						    spci_vcpu_index_t vcpu_idx)
 {
 	return hf_vcpu_run_return_decode(
 		hf_call(HF_VCPU_RUN, vm_id, vcpu_idx, 0));
@@ -218,7 +218,7 @@
  *    up or kick the target vCPU.
  */
 static inline int64_t hf_interrupt_inject(spci_vm_id_t target_vm_id,
-					  uint32_t target_vcpu_idx,
+					  spci_vcpu_index_t target_vcpu_idx,
 					  uint32_t intid)
 {
 	return hf_call(HF_INTERRUPT_INJECT, target_vm_id, target_vcpu_idx,
diff --git a/inc/vmapi/hf/spci.h b/inc/vmapi/hf/spci.h
index c387e62..25e8e83 100644
--- a/inc/vmapi/hf/spci.h
+++ b/inc/vmapi/hf/spci.h
@@ -68,6 +68,7 @@
  * different name to make the different semantics clear.
  */
 typedef spci_vm_id_t spci_vm_count_t;
+typedef uint16_t spci_vcpu_index_t;
 
 /** SPCI common message header. */
 struct spci_message {
diff --git a/src/api.c b/src/api.c
index ccbeeda..91b582d 100644
--- a/src/api.c
+++ b/src/api.c
@@ -22,6 +22,7 @@
 #include "hf/assert.h"
 #include "hf/dlog.h"
 #include "hf/mm.h"
+#include "hf/spci.h"
 #include "hf/spinlock.h"
 #include "hf/std.h"
 #include "hf/vm.h"
@@ -481,7 +482,8 @@
 /**
  * Runs the given vcpu of the given vm.
  */
-struct hf_vcpu_run_return api_vcpu_run(spci_vm_id_t vm_id, uint32_t vcpu_idx,
+struct hf_vcpu_run_return api_vcpu_run(spci_vm_id_t vm_id,
+				       spci_vcpu_index_t vcpu_idx,
 				       const struct vcpu *current,
 				       struct vcpu **next)
 {
@@ -1159,7 +1161,7 @@
  *    up or kick the target vCPU.
  */
 int64_t api_interrupt_inject(spci_vm_id_t target_vm_id,
-			     uint32_t target_vcpu_idx, uint32_t intid,
+			     spci_vcpu_index_t target_vcpu_idx, uint32_t intid,
 			     struct vcpu *current, struct vcpu **next)
 {
 	struct vcpu *target_vcpu;
diff --git a/src/arch/aarch64/hypervisor/handler.c b/src/arch/aarch64/hypervisor/handler.c
index 7943401..04223db 100644
--- a/src/arch/aarch64/hypervisor/handler.c
+++ b/src/arch/aarch64/hypervisor/handler.c
@@ -134,7 +134,7 @@
 void maybe_invalidate_tlb(struct vcpu *vcpu)
 {
 	size_t current_cpu_index = cpu_index(vcpu->cpu);
-	size_t new_vcpu_index = vcpu_index(vcpu);
+	spci_vcpu_index_t new_vcpu_index = vcpu_index(vcpu);
 
 	if (vcpu->vm->arch.last_vcpu_on_cpu[current_cpu_index] !=
 	    new_vcpu_index) {
diff --git a/src/arch/aarch64/hypervisor/psci_handler.c b/src/arch/aarch64/hypervisor/psci_handler.c
index 76e4798..2736814 100644
--- a/src/arch/aarch64/hypervisor/psci_handler.c
+++ b/src/arch/aarch64/hypervisor/psci_handler.c
@@ -24,6 +24,7 @@
 #include "hf/cpu.h"
 #include "hf/dlog.h"
 #include "hf/panic.h"
+#include "hf/spci.h"
 #include "hf/vm.h"
 
 #include "psci.h"
@@ -227,7 +228,7 @@
  * Convert a PSCI CPU / affinity ID for a secondary VM to the corresponding vCPU
  * index.
  */
-uint32_t vcpu_id_to_index(uint64_t vcpu_id)
+spci_vcpu_index_t vcpu_id_to_index(uint64_t vcpu_id)
 {
 	/* For now we use indices as IDs for the purposes of PSCI. */
 	return vcpu_id;
@@ -282,7 +283,8 @@
 		uint32_t lowest_affinity_level = arg1;
 		struct vm *vm = vcpu->vm;
 		struct vcpu_locked target_vcpu;
-		uint32_t target_vcpu_index = vcpu_id_to_index(target_affinity);
+		spci_vcpu_index_t target_vcpu_index =
+			vcpu_id_to_index(target_affinity);
 
 		if (lowest_affinity_level != 0) {
 			/* Affinity levels greater than 0 not supported. */
@@ -327,7 +329,8 @@
 		uint64_t target_cpu = arg0;
 		ipaddr_t entry_point_address = ipa_init(arg1);
 		uint64_t context_id = arg2;
-		uint32_t target_vcpu_index = vcpu_id_to_index(target_cpu);
+		spci_vcpu_index_t target_vcpu_index =
+			vcpu_id_to_index(target_cpu);
 		struct vm *vm = vcpu->vm;
 		struct vcpu *target_vcpu;
 
diff --git a/src/arch/aarch64/inc/hf/arch/types.h b/src/arch/aarch64/inc/hf/arch/types.h
index cc1cbef..527ce92 100644
--- a/src/arch/aarch64/inc/hf/arch/types.h
+++ b/src/arch/aarch64/inc/hf/arch/types.h
@@ -20,6 +20,7 @@
 #include <stdint.h>
 
 #include "hf/assert.h"
+#include "hf/spci.h"
 
 #define PAGE_BITS 12
 #define PAGE_LEVEL_BITS 9
@@ -58,7 +59,7 @@
 	 * on that CPU, which avoids contention and so no lock is needed to
 	 * access this field.
 	 */
-	uint32_t last_vcpu_on_cpu[MAX_CPUS];
+	spci_vcpu_index_t last_vcpu_on_cpu[MAX_CPUS];
 };
 
 /** Type to represent the register state of a vCPU.  */
diff --git a/src/cpu.c b/src/cpu.c
index affba4a..eeea904 100644
--- a/src/cpu.c
+++ b/src/cpu.c
@@ -22,6 +22,7 @@
 
 #include "hf/api.h"
 #include "hf/dlog.h"
+#include "hf/spci.h"
 #include "hf/std.h"
 #include "hf/vm.h"
 
@@ -203,9 +204,12 @@
 	vcpu.vcpu->state = VCPU_STATE_READY;
 }
 
-size_t vcpu_index(const struct vcpu *vcpu)
+spci_vcpu_index_t vcpu_index(const struct vcpu *vcpu)
 {
-	return vcpu - vcpu->vm->vcpus;
+	size_t index = vcpu - vcpu->vm->vcpus;
+
+	assert(index < UINT16_MAX);
+	return index;
 }
 
 /**
diff --git a/src/vm.c b/src/vm.c
index 8532144..7c5166b 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -18,6 +18,7 @@
 
 #include "hf/api.h"
 #include "hf/cpu.h"
+#include "hf/spci.h"
 #include "hf/std.h"
 
 #include "vmapi/hf/call.h"
@@ -112,7 +113,7 @@
  * Get the vCPU with the given index from the given VM.
  * This assumes the index is valid, i.e. less than vm->vcpu_count.
  */
-struct vcpu *vm_get_vcpu(struct vm *vm, uint32_t vcpu_index)
+struct vcpu *vm_get_vcpu(struct vm *vm, spci_vcpu_index_t vcpu_index)
 {
 	assert(vcpu_index < vm->vcpu_count);
 	return &vm->vcpus[vcpu_index];