mm_context_id_t id;
        u16 user_psize;         /* page size index */
 
+       /* Number of bits in the mm_cpumask */
+       atomic_t active_cpus;
+
        /* NPU NMMU context */
        struct npu_context *npu_context;
 
 
                                   struct mm_struct *mm) { }
 #endif
 
+#ifdef CONFIG_PPC_BOOK3S_64
+static inline void inc_mm_active_cpus(struct mm_struct *mm)
+{
+       atomic_inc(&mm->context.active_cpus);
+}
+#else
+static inline void inc_mm_active_cpus(struct mm_struct *mm) { }
+#endif
 
 /*
  * switch_mm is the entry point called from the architecture independent
        /* Mark this context has been used on the new CPU */
        if (!cpumask_test_cpu(smp_processor_id(), mm_cpumask(next))) {
                cpumask_set_cpu(smp_processor_id(), mm_cpumask(next));
+               inc_mm_active_cpus(next);
 
                /*
                 * This full barrier orders the store to the cpumask above vs
 
                              topology_sibling_cpumask(smp_processor_id()));
 }
 
+#ifdef CONFIG_PPC_BOOK3S_64
+static inline int mm_is_thread_local(struct mm_struct *mm)
+{
+       if (atomic_read(&mm->context.active_cpus) > 1)
+               return false;
+       return cpumask_test_cpu(smp_processor_id(), mm_cpumask(mm));
+}
+#else /* CONFIG_PPC_BOOK3S_64 */
 static inline int mm_is_thread_local(struct mm_struct *mm)
 {
        return cpumask_equal(mm_cpumask(mm),
                              cpumask_of(smp_processor_id()));
 }
+#endif /* !CONFIG_PPC_BOOK3S_64 */
 
-#else
+#else /* CONFIG_SMP */
 static inline int mm_is_core_local(struct mm_struct *mm)
 {
        return 1;
 
 #ifdef CONFIG_SPAPR_TCE_IOMMU
        mm_iommu_init(mm);
 #endif
+       atomic_set(&mm->context.active_cpus, 0);
+
        return 0;
 }