vcpu pinning

Signed-off-by: Oliver Upton <oliver.upton@linux.dev>
diff --git a/arm/aarch64/include/kvm/kvm-config-arch.h b/arm/aarch64/include/kvm/kvm-config-arch.h
index 1c40bba..74fc0fc 100644
--- a/arm/aarch64/include/kvm/kvm-config-arch.h
+++ b/arm/aarch64/include/kvm/kvm-config-arch.h
@@ -15,6 +15,10 @@
 	OPT_CALLBACK('\0', "vcpu-affinity", kvm, "cpulist",  		\
 			"Specify the CPU affinity that will apply to "	\
 			"all VCPUs", vcpu_affinity_parser, kvm),	\
+	OPT_BOOLEAN('\0', "vcpu-pinned", &(cfg)->vcpu_pinned,		\
+			"1:1 pin vCPUs to CPUs in the vcpu-affinity"	\
+			"list. You must specify at least as many CPUs"	\
+			"as vCPUs in the vcpu-affinity list."),		\
 	OPT_U64('\0', "kaslr-seed", &(cfg)->kaslr_seed,			\
 			"Specify random seed for Kernel Address Space "	\
 			"Layout Randomization (KASLR)"),		\
diff --git a/arm/aarch64/kvm-cpu.c b/arm/aarch64/kvm-cpu.c
index 316e20c..fd97833 100644
--- a/arm/aarch64/kvm-cpu.c
+++ b/arm/aarch64/kvm-cpu.c
@@ -153,17 +153,53 @@
 	return 0;
 }
 
+static cpu_set_t *kvm_cpu__get_affinity(struct kvm_cpu *vcpu)
+{
+	size_t setsize = CPU_ALLOC_SIZE(NR_CPUS);
+	cpu_set_t *affinity = CPU_ALLOC(NR_CPUS);
+	struct kvm *kvm = vcpu->kvm;
+	unsigned long nr_cpus;
+	int i;
+
+	if (!kvm->arch.vcpu_affinity_cpuset || !affinity)
+		return NULL;
+
+	CPU_ZERO_S(CPU_ALLOC_SIZE(NR_CPUS), affinity);
+
+	if (!kvm->cfg.arch.vcpu_pinned) {
+		CPU_OR_S(setsize, affinity, affinity, kvm->arch.vcpu_affinity_cpuset);
+		return affinity;
+	}
+
+	for (i = 0, nr_cpus = 0; i < NR_CPUS; i++) {
+		if (!CPU_ISSET_S(i, setsize, kvm->arch.vcpu_affinity_cpuset))
+			continue;
+
+		if (nr_cpus == vcpu->cpu_id) {
+			CPU_SET_S(i, setsize, affinity);
+			return affinity;
+		}
+
+		nr_cpus++;
+	}
+
+	die("Invalid vcpu-affinity mask for vcpu pinning");
+	return NULL;
+}
+
 void kvm_cpu__reset_vcpu(struct kvm_cpu *vcpu)
 {
 	struct kvm *kvm = vcpu->kvm;
 	cpu_set_t *affinity;
 	int ret;
 
-	affinity = kvm->arch.vcpu_affinity_cpuset;
+	affinity = kvm_cpu__get_affinity(vcpu);
 	if (affinity) {
 		ret = sched_setaffinity(0, sizeof(cpu_set_t), affinity);
 		if (ret == -1)
 			die_perror("sched_setaffinity");
+
+		CPU_FREE(affinity);
 	}
 
 	if (kvm->cfg.arch.aarch32_guest)
diff --git a/arm/include/arm-common/kvm-config-arch.h b/arm/include/arm-common/kvm-config-arch.h
index 223b5a9..b35eeeb 100644
--- a/arm/include/arm-common/kvm-config-arch.h
+++ b/arm/include/arm-common/kvm-config-arch.h
@@ -15,6 +15,7 @@
 	u64		fw_addr;
 	bool no_pvtime;
 	bool		in_kernel_smccc;
+	bool		vcpu_pinned;
 };
 
 int irqchip_parser(const struct option *opt, const char *arg, int unset);