Now that unuse_temporary_mm() lives in tlb.c it can access loaded_mm.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 arch/x86/include/asm/mmu_context.h |    2 +-
 arch/x86/kernel/alternative.c      |    2 +-
 arch/x86/mm/tlb.c                  |    8 +++++---
 3 files changed, 7 insertions(+), 5 deletions(-)

--- a/arch/x86/include/asm/mmu_context.h
+++ b/arch/x86/include/asm/mmu_context.h
@@ -268,6 +268,6 @@ typedef struct {
 } temp_mm_state_t;
 
 extern temp_mm_state_t use_temporary_mm(struct mm_struct *mm);
-extern void unuse_temporary_mm(struct mm_struct *mm, temp_mm_state_t prev_state);
+extern void unuse_temporary_mm(temp_mm_state_t prev_state);
 
 #endif /* _ASM_X86_MMU_CONTEXT_H */
--- a/arch/x86/kernel/alternative.c
+++ b/arch/x86/kernel/alternative.c
@@ -1874,7 +1874,7 @@ static void *__text_poke(text_poke_f fun
 	 * instruction that already allows the core to see the updated version.
 	 * Xen-PV is assumed to serialize execution in a similar manner.
 	 */
-	unuse_temporary_mm(poking_mm, prev);
+	unuse_temporary_mm(prev);
 
 	/*
 	 * Flushing the TLB might involve IPIs, which would require enabled
--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -720,13 +720,15 @@ temp_mm_state_t use_temporary_mm(struct
 	return temp_state;
 }
 
-void unuse_temporary_mm(struct mm_struct *mm, temp_mm_state_t prev_state)
+void unuse_temporary_mm(temp_mm_state_t prev_state)
 {
 	lockdep_assert_irqs_disabled();
-	switch_mm_irqs_off(NULL, prev_state.mm, current);
 
 	/* Clear the cpumask, to indicate no TLB flushing is needed anywhere */
-	cpumask_clear_cpu(raw_smp_processor_id(), mm_cpumask(mm));
+	cpumask_clear_cpu(smp_processor_id(),
+			  mm_cpumask(this_cpu_read(cpu_tlbstate.loaded_mm)));
+
+	switch_mm_irqs_off(NULL, prev_state.mm, current);
 
 	/*
 	 * Restore the breakpoints if they were disabled before the temporary mm