powerpc/64s/radix: Fix preempt imbalance in TLB flush

Fixes: 424de9c6e3f8 ("powerpc/mm/radix: Avoid flushing the PWC on every flush_tlb_range")
Signed-off-by: Nicholas Piggin <npiggin@gmail.com>
Signed-off-by: Michael Ellerman <mpe@ellerman.id.au>

1 files changed, 2 insertions, 0 deletions

diff --git a/arch/powerpc/mm/tlb-radix.c b/arch/powerpc/mm/tlb-radix.c
index b3e849c4886e..d304028641a2 100644
--- a/arch/powerpc/mm/tlb-radix.c
+++ b/arch/powerpc/mm/tlb-radix.c
@@ -360,12 +360,14 @@ void radix__flush_tlb_collapsed_pmd(struct mm_struct *mm, unsigned long addr)
 
 
         pid = mm ? mm->context.id : 0;
+        preempt_disable();
         if (unlikely(pid == MMU_NO_CONTEXT))
                 goto no_context;
 
         /* 4k page size, just blow the world */
         if (PAGE_SIZE == 0x1000) {
                 radix__flush_all_mm(mm);
+                preempt_enable();
                 return;
         }