--- kvm-28/kernel/vmx.c	2007-06-10 11:54:27.000000000 +0200
+++ kvm-28-preempt/kernel/vmx.c	2007-06-20 12:59:57.000000000 +0200
@@ -375,10 +375,16 @@
 static void vmx_vcpu_load(struct kvm_vcpu *vcpu)
 {
 	u64 phys_addr = __pa(vcpu->vmcs);
-	int cpu;
-
-	cpu = get_cpu();
-
+	int cpu = raw_smp_processor_id();
+	cpumask_t this_mask = cpumask_of_cpu(cpu);
+ 
+	/*
+	 * Keep the context preemptible, but do not migrate
+	 * away to another CPU. TODO: make sure this persists.
+	 * Save/restore original mask.
+	 */
+	if (unlikely(!cpus_equal(current->cpus_allowed, this_mask)))
+		set_cpus_allowed(current, cpumask_of_cpu(cpu));
 	if (vcpu->cpu != cpu)
 		vcpu_clear(vcpu);
 
@@ -416,7 +422,6 @@
 {
 	vmx_load_host_state(vcpu);
 	kvm_put_guest_fpu(vcpu);
-	put_cpu();
 }
 
 static void vmx_fpu_activate(struct kvm_vcpu *vcpu)
@@ -441,7 +446,9 @@
 
 static void vmx_vcpu_decache(struct kvm_vcpu *vcpu)
 {
+	preempt_disable();
 	vcpu_clear(vcpu);
+	preempt_enable();
 }
 
 static unsigned long vmx_get_rflags(struct kvm_vcpu *vcpu)