mp->vmaddr = vmaddr & PMD_MASK;
        INIT_LIST_HEAD(&mp->mapper);
        page->index = (unsigned long) mp;
-       atomic_set(&page->_mapcount, 3);
+       atomic_set(&page->_mapcount, 0);
        table = (unsigned long *) page_to_phys(page);
 -      clear_table(table, _PAGE_TYPE_EMPTY, PAGE_SIZE/2);
 -      clear_table(table + PTRS_PER_PTE, 0, PAGE_SIZE/2);
 +      clear_table(table, _PAGE_INVALID, PAGE_SIZE/2);
 +      clear_table(table + PTRS_PER_PTE, PGSTE_HR_BIT | PGSTE_HC_BIT,
 +                  PAGE_SIZE/2);
        return table;
  }
  
        if (mm_has_pgste(tsk->mm))
                return 0;
  
-       /* lets check if we are allowed to replace the mm */
-       task_lock(tsk);
-       if (!tsk->mm || atomic_read(&tsk->mm->mm_users) > 1 ||
- #ifdef CONFIG_AIO
-           !hlist_empty(&tsk->mm->ioctx_list) ||
- #endif
-           tsk->mm != tsk->active_mm) {
-               task_unlock(tsk);
-               return -EINVAL;
-       }
-       task_unlock(tsk);
- 
-       /* we copy the mm and let dup_mm create the page tables with_pgstes */
-       tsk->mm->context.alloc_pgste = 1;
-       /* make sure that both mms have a correct rss state */
-       sync_mm_rss(tsk->mm);
-       mm = dup_mm(tsk);
-       tsk->mm->context.alloc_pgste = 0;
-       if (!mm)
-               return -ENOMEM;
- 
- #ifdef CONFIG_TRANSPARENT_HUGEPAGE
+       down_write(&mm->mmap_sem);
        /* split thp mappings and disable thp for future mappings */
        thp_split_mm(mm);
-       mm->def_flags |= VM_NOHUGEPAGE;
- #endif
- 
-       /* Now lets check again if something happened */
-       task_lock(tsk);
-       if (!tsk->mm || atomic_read(&tsk->mm->mm_users) > 1 ||
- #ifdef CONFIG_AIO
-           !hlist_empty(&tsk->mm->ioctx_list) ||
- #endif
-           tsk->mm != tsk->active_mm) {
-               mmput(mm);
-               task_unlock(tsk);
-               return -EINVAL;
-       }
- 
-       /* ok, we are alone. No ptrace, no threads, etc. */
-       old_mm = tsk->mm;
-       tsk->mm = tsk->active_mm = mm;
-       preempt_disable();
-       update_mm(mm, tsk);
-       atomic_inc(&mm->context.attach_count);
-       atomic_dec(&old_mm->context.attach_count);
-       cpumask_set_cpu(smp_processor_id(), mm_cpumask(mm));
-       preempt_enable();
-       task_unlock(tsk);
-       mmput(old_mm);
-       return 0;
+       /* Reallocate the page tables with pgstes */
+       mm->context.has_pgste = 1;
 -      tlb_gather_mmu(&tlb, mm, 0);
++      tlb_gather_mmu(&tlb, mm, 0, TASK_SIZE);
+       page_table_realloc(&tlb, mm, 0, TASK_SIZE);
 -      tlb_finish_mmu(&tlb, 0, -1);
++      tlb_finish_mmu(&tlb, 0, TASK_SIZE);
+       up_write(&mm->mmap_sem);
+       return mm->context.has_pgste ? 0 : -ENOMEM;
  }
  EXPORT_SYMBOL_GPL(s390_enable_sie);