8 files changed, 28 insertions, 6 deletions

diff --git a/arch/s390/include/asm/hugetlb.h b/arch/s390/include/asm/hugetlb.h
index 670a1d1745d2..bb8343d157bc 100644
--- a/arch/s390/include/asm/hugetlb.h
+++ b/arch/s390/include/asm/hugetlb.h
@@ -97,6 +97,7 @@ static inline pte_t huge_ptep_get_and_clear(struct mm_struct *mm,
 {
         pte_t pte = huge_ptep_get(ptep);
 
+        mm->context.flush_mm = 1;
         pmd_clear((pmd_t *) ptep);
         return pte;
 }
@@ -167,7 +168,8 @@ static inline void huge_ptep_invalidate(struct mm_struct *mm,
 ({                                                                      \
         pte_t __pte = huge_ptep_get(__ptep);                            \
         if (pte_write(__pte)) {                                         \
-                if (atomic_read(&(__mm)->mm_users) > 1 ||               \
+                (__mm)->context.flush_mm = 1;                           \
+                if (atomic_read(&(__mm)->context.attach_count) > 1 ||   \
                     (__mm) != current->active_mm)                       \
                         huge_ptep_invalidate(__mm, __addr, __ptep);     \
                 set_huge_pte_at(__mm, __addr, __ptep,                   \
diff --git a/arch/s390/include/asm/mmu.h b/arch/s390/include/asm/mmu.h
index 99e3409102b9..78522cdefdd4 100644
--- a/arch/s390/include/asm/mmu.h
+++ b/arch/s390/include/asm/mmu.h
@@ -2,6 +2,8 @@
 #define __MMU_H
 
 typedef struct {
+        atomic_t attach_count;
+        unsigned int flush_mm;
         spinlock_t list_lock;
         struct list_head crst_list;
         struct list_head pgtable_list;
diff --git a/arch/s390/include/asm/mmu_context.h b/arch/s390/include/asm/mmu_context.h
index 976e273988c2..a6f0e7cc9cde 100644
--- a/arch/s390/include/asm/mmu_context.h
+++ b/arch/s390/include/asm/mmu_context.h
@@ -11,11 +11,14 @@
 
 #include <asm/pgalloc.h>
 #include <asm/uaccess.h>
+#include <asm/tlbflush.h>
 #include <asm-generic/mm_hooks.h>
 
 static inline int init_new_context(struct task_struct *tsk,
                                    struct mm_struct *mm)
 {
+        atomic_set(&mm->context.attach_count, 0);
+        mm->context.flush_mm = 0;
         mm->context.asce_bits = _ASCE_TABLE_LENGTH | _ASCE_USER_BITS;
 #ifdef CONFIG_64BIT
         mm->context.asce_bits |= _ASCE_TYPE_REGION3;
@@ -76,6 +79,12 @@ static inline void switch_mm(struct mm_struct *prev, struct mm_struct *next,
 {
         cpumask_set_cpu(smp_processor_id(), mm_cpumask(next));
         update_mm(next, tsk);
+        atomic_dec(&prev->context.attach_count);
+        WARN_ON(atomic_read(&prev->context.attach_count) < 0);
+        atomic_inc(&next->context.attach_count);
+        /* Check for TLBs not flushed yet */
+        if (next->context.flush_mm)
+                __tlb_flush_mm(next);
 }
 
 #define enter_lazy_tlb(mm,tsk)  do { } while (0)
diff --git a/arch/s390/include/asm/pgtable.h b/arch/s390/include/asm/pgtable.h
index 89a504c3f12e..3157441ee1da 100644
--- a/arch/s390/include/asm/pgtable.h
+++ b/arch/s390/include/asm/pgtable.h
@@ -880,7 +880,8 @@ static inline void ptep_invalidate(struct mm_struct *mm,
 #define ptep_get_and_clear(__mm, __address, __ptep)                     \
 ({                                                                      \
         pte_t __pte = *(__ptep);                                        \
-        if (atomic_read(&(__mm)->mm_users) > 1 ||                       \
+        (__mm)->context.flush_mm = 1;                                   \
+        if (atomic_read(&(__mm)->context.attach_count) > 1 ||           \
             (__mm) != current->active_mm)                               \
                 ptep_invalidate(__mm, __address, __ptep);               \
         else                                                            \
@@ -923,7 +924,8 @@ static inline pte_t ptep_get_and_clear_full(struct mm_struct *mm,
 ({                                                                      \
         pte_t __pte = *(__ptep);                                        \
         if (pte_write(__pte)) {                                         \
-                if (atomic_read(&(__mm)->mm_users) > 1 ||               \
+                (__mm)->context.flush_mm = 1;                           \
+                if (atomic_read(&(__mm)->context.attach_count) > 1 ||   \
                     (__mm) != current->active_mm)                       \
                         ptep_invalidate(__mm, __addr, __ptep);          \
                 set_pte_at(__mm, __addr, __ptep, pte_wrprotect(__pte)); \
diff --git a/arch/s390/include/asm/tlb.h b/arch/s390/include/asm/tlb.h
index 81150b053689..fd1c00d08bf5 100644
--- a/arch/s390/include/asm/tlb.h
+++ b/arch/s390/include/asm/tlb.h
@@ -50,8 +50,7 @@ static inline struct mmu_gather *tlb_gather_mmu(struct mm_struct *mm,
         struct mmu_gather *tlb = &get_cpu_var(mmu_gathers);
 
         tlb->mm = mm;
-        tlb->fullmm = full_mm_flush || (num_online_cpus() == 1) ||
-                (atomic_read(&mm->mm_users) <= 1 && mm == current->active_mm);
+        tlb->fullmm = full_mm_flush;
         tlb->nr_ptes = 0;
         tlb->nr_pxds = TLB_NR_PTRS;
         if (tlb->fullmm)
diff --git a/arch/s390/include/asm/tlbflush.h b/arch/s390/include/asm/tlbflush.h
index 304cffa623e1..29d5d6d4becc 100644
--- a/arch/s390/include/asm/tlbflush.h
+++ b/arch/s390/include/asm/tlbflush.h
@@ -94,8 +94,12 @@ static inline void __tlb_flush_mm(struct mm_struct * mm)
 
 static inline void __tlb_flush_mm_cond(struct mm_struct * mm)
 {
-        if (atomic_read(&mm->mm_users) <= 1 && mm == current->active_mm)
+        spin_lock(&mm->page_table_lock);
+        if (mm->context.flush_mm) {
                 __tlb_flush_mm(mm);
+                mm->context.flush_mm = 0;
+        }
+        spin_unlock(&mm->page_table_lock);
 }
 
 /*
diff --git a/arch/s390/kernel/smp.c b/arch/s390/kernel/smp.c
index 541053ed234e..8127ebd59c4d 100644
--- a/arch/s390/kernel/smp.c
+++ b/arch/s390/kernel/smp.c
@@ -583,6 +583,7 @@ int __cpuinit __cpu_up(unsigned int cpu)
         sf->gprs[9] = (unsigned long) sf;
         cpu_lowcore->save_area[15] = (unsigned long) sf;
         __ctl_store(cpu_lowcore->cregs_save_area, 0, 15);
+        atomic_inc(&init_mm.context.attach_count);
         asm volatile(
                 "       stam    0,15,0(%0)"
                 : : "a" (&cpu_lowcore->access_regs_save_area) : "memory");
@@ -659,6 +660,7 @@ void __cpu_die(unsigned int cpu)
         while (sigp_p(0, cpu, sigp_set_prefix) == sigp_busy)
                 udelay(10);
         smp_free_lowcore(cpu);
+        atomic_dec(&init_mm.context.attach_count);
         pr_info("Processor %d stopped\n", cpu);
 }
 
diff --git a/arch/s390/mm/init.c b/arch/s390/mm/init.c
index acc91c75bc94..30eb6d02ddb8 100644
--- a/arch/s390/mm/init.c
+++ b/arch/s390/mm/init.c
@@ -74,6 +74,8 @@ void __init paging_init(void)
         __ctl_load(S390_lowcore.kernel_asce, 13, 13);
         __raw_local_irq_ssm(ssm_mask);
 
+        atomic_set(&init_mm.context.attach_count, 1);
+
         sparse_memory_present_with_active_regions(MAX_NUMNODES);
         sparse_init();
         memset(max_zone_pfns, 0, sizeof(max_zone_pfns));