switch_mm race condition with Ingo's scheduler

diff -urNp linux-2.5.18-ia64/arch/ia64/mm/tlb.c linux-2.5.18-ia64-wrap/arch/ia64/mm/tlb.c
@@ -77,7 +77,11 @@ wrap_mmu_context (struct mm_struct *mm)
 			ia64_ctx.limit = tsk_context;
 	}
 	read_unlock(&tasklist_lock);
-	flush_tlb_all();
+	//flush_tlb_all(); /* potential race condition with O(1) scheduler [EF] */
+	for (i=0; i<smp_num_cpus; i++)
+		cpu_data(i)->tlb_flush = 1;
+	__flush_tlb_all();
+	local_cpu_data->tlb_flush = 0;
 }
 
 void

diff -urNp linux-2.5.18-ia64/include/asm-ia64/mmu_context.h linux-2.5.18-ia64-wrap/include/asm-ia64/mmu_context.h
@@ -44,6 +44,23 @@ enter_lazy_tlb (struct mm_struct *mm, st
 {
 }
 
+/* 
+ * When the context counter wraps around all TLBs need to be flushed because
+ * an old context number might have been reused. This is signalled by a bit
+ * set in ia64_ctx.flush, which is checked in the routine below. Called by
+ * activate_mm().                                        <[email protected]>
+ */
+static inline void
+delayed_tlb_flush (void)
+{
+	extern void __flush_tlb_all (void);
+
+	if (unlikely(local_cpu_data->tlb_flush)) {
+		__flush_tlb_all();
+		local_cpu_data->tlb_flush = 0;
+	}
+}
+
 static inline void
 get_new_mmu_context (struct mm_struct *mm)
 {

@@ -54,7 +71,6 @@ get_new_mmu_context (struct mm_struct *m
 		mm->context = ia64_ctx.next++;
 	}
 	spin_unlock(&ia64_ctx.lock);
-
 }
 
 static inline void

@@ -113,6 +129,7 @@ activate_mm (struct mm_struct *prev, str
 	 * We may get interrupts here, but that's OK because interrupt
 	 * handlers cannot touch user-space.
 	 */
+	delayed_tlb_flush();
 	ia64_set_kr(IA64_KR_PT_BASE, __pa(next->pgd));
 	get_mmu_context(next);
 	reload_context(next);

diff -urNp linux-2.5.18-ia64/include/asm-ia64/processor.h linux-2.5.18-ia64-wrap/include/asm-ia64/processor.h
@@ -151,6 +151,7 @@ extern struct cpuinfo_ia64 {
 	/* CPUID-derived information: */
 	__u64 ppn;
 	__u64 features;
+	__u8 tlb_flush;		/* flush TLB before next context switch if non-zero */
 	__u8 number;
 	__u8 revision;
 	__u8 model;