Merge branch 'next' of git://git.kernel.org/pub/scm/virt/kvm/kvm

Pull KVM updates from Gleb Natapov:
 "The highlights of the release are nested EPT and pv-ticketlocks
  support (hypervisor part, guest part, which is most of the code, goes
  through tip tree).  Apart of that there are many fixes for all arches"

Fix up semantic conflicts as discussed in the pull request thread..

* 'next' of git://git.kernel.org/pub/scm/virt/kvm/kvm: (88 commits)
  ARM: KVM: Add newlines to panic strings
  ARM: KVM: Work around older compiler bug
  ARM: KVM: Simplify tracepoint text
  ARM: KVM: Fix kvm_set_pte assignment
  ARM: KVM: vgic: Bump VGIC_NR_IRQS to 256
  ARM: KVM: Bugfix: vgic_bytemap_get_reg per cpu regs
  ARM: KVM: vgic: fix GICD_ICFGRn access
  ARM: KVM: vgic: simplify vgic_get_target_reg
  KVM: MMU: remove unused parameter
  KVM: PPC: Book3S PR: Rework kvmppc_mmu_book3s_64_xlate()
  KVM: PPC: Book3S PR: Make instruction fetch fallback work for system calls
  KVM: PPC: Book3S PR: Don't corrupt guest state when kernel uses VMX
  KVM: x86: update masterclock when kvmclock_offset is calculated (v2)
  KVM: PPC: Book3S: Fix compile error in XICS emulation
  KVM: PPC: Book3S PR: return appropriate error when allocation fails
  arch: powerpc: kvm: add signed type cast for comparation
  KVM: x86: add comments where MMIO does not return to the emulator
  KVM: vmx: count exits to userspace during invalid guest emulation
  KVM: rename __kvm_io_bus_sort_cmp to kvm_io_bus_cmp
  kvm: optimize away THP checks in kvm_is_mmio_pfn()
  ...

10 files changed, 195 insertions, 116 deletions

diff --git a/arch/s390/include/asm/kvm_host.h b/arch/s390/include/asm/kvm_host.h
index 3238d4004e84..e87ecaa2c569 100644
--- a/arch/s390/include/asm/kvm_host.h
+++ b/arch/s390/include/asm/kvm_host.h
@@ -274,6 +274,14 @@ struct kvm_arch{
         int css_support;
 };
 
+#define KVM_HVA_ERR_BAD         (-1UL)
+#define KVM_HVA_ERR_RO_BAD      (-2UL)
+
+static inline bool kvm_is_error_hva(unsigned long addr)
+{
+        return IS_ERR_VALUE(addr);
+}
+
 extern int sie64a(struct kvm_s390_sie_block *, u64 *);
 extern char sie_exit;
 #endif
diff --git a/arch/s390/include/asm/mmu.h b/arch/s390/include/asm/mmu.h
index 6340178748bf..ff132ac64ddd 100644
--- a/arch/s390/include/asm/mmu.h
+++ b/arch/s390/include/asm/mmu.h
@@ -12,8 +12,6 @@ typedef struct {
         unsigned long asce_bits;
         unsigned long asce_limit;
         unsigned long vdso_base;
-        /* Cloned contexts will be created with extended page tables. */
-        unsigned int alloc_pgste:1;
         /* The mmu context has extended page tables. */
         unsigned int has_pgste:1;
 } mm_context_t;
diff --git a/arch/s390/include/asm/mmu_context.h b/arch/s390/include/asm/mmu_context.h
index 7b7fce4e8469..9f973d8de90e 100644
--- a/arch/s390/include/asm/mmu_context.h
+++ b/arch/s390/include/asm/mmu_context.h
@@ -21,24 +21,7 @@ static inline int init_new_context(struct task_struct *tsk,
 #ifdef CONFIG_64BIT
         mm->context.asce_bits |= _ASCE_TYPE_REGION3;
 #endif
-        if (current->mm && current->mm->context.alloc_pgste) {
-                /*
-                 * alloc_pgste indicates, that any NEW context will be created
-                 * with extended page tables. The old context is unchanged. The
-                 * page table allocation and the page table operations will
-                 * look at has_pgste to distinguish normal and extended page
-                 * tables. The only way to create extended page tables is to
-                 * set alloc_pgste and then create a new context (e.g. dup_mm).
-                 * The page table allocation is called after init_new_context
-                 * and if has_pgste is set, it will create extended page
-                 * tables.
-                 */
-                mm->context.has_pgste = 1;
-                mm->context.alloc_pgste = 1;
-        } else {
-                mm->context.has_pgste = 0;
-                mm->context.alloc_pgste = 0;
-        }
+        mm->context.has_pgste = 0;
         mm->context.asce_limit = STACK_TOP_MAX;
         crst_table_init((unsigned long *) mm->pgd, pgd_entry_type(mm));
         return 0;
diff --git a/arch/s390/include/asm/pgtable.h b/arch/s390/include/asm/pgtable.h
index 9f215b40109e..9b60a36c348d 100644
--- a/arch/s390/include/asm/pgtable.h
+++ b/arch/s390/include/asm/pgtable.h
@@ -1442,6 +1442,17 @@ static inline pmd_t pmd_mkwrite(pmd_t pmd)
 }
 #endif /* CONFIG_TRANSPARENT_HUGEPAGE || CONFIG_HUGETLB_PAGE */
 
+static inline void pmdp_flush_lazy(struct mm_struct *mm,
+                                   unsigned long address, pmd_t *pmdp)
+{
+        int active = (mm == current->active_mm) ? 1 : 0;
+
+        if ((atomic_read(&mm->context.attach_count) & 0xffff) > active)
+                __pmd_idte(address, pmdp);
+        else
+                mm->context.flush_mm = 1;
+}
+
 #ifdef CONFIG_TRANSPARENT_HUGEPAGE
 
 #define __HAVE_ARCH_PGTABLE_DEPOSIT
diff --git a/arch/s390/include/asm/processor.h b/arch/s390/include/asm/processor.h
index b0e6435b2f02..0eb37505cab1 100644
--- a/arch/s390/include/asm/processor.h
+++ b/arch/s390/include/asm/processor.h
@@ -43,6 +43,7 @@ extern void execve_tail(void);
 #ifndef CONFIG_64BIT
 
 #define TASK_SIZE               (1UL << 31)
+#define TASK_MAX_SIZE           (1UL << 31)
 #define TASK_UNMAPPED_BASE      (1UL << 30)
 
 #else /* CONFIG_64BIT */
@@ -51,6 +52,7 @@ extern void execve_tail(void);
 #define TASK_UNMAPPED_BASE      (test_thread_flag(TIF_31BIT) ? \
                                         (1UL << 30) : (1UL << 41))
 #define TASK_SIZE               TASK_SIZE_OF(current)
+#define TASK_MAX_SIZE           (1UL << 53)
 
 #endif /* CONFIG_64BIT */
 
diff --git a/arch/s390/kvm/diag.c b/arch/s390/kvm/diag.c
index 3074475c8ae0..3a74d8af0d69 100644
--- a/arch/s390/kvm/diag.c
+++ b/arch/s390/kvm/diag.c
@@ -119,12 +119,21 @@ static int __diag_virtio_hypercall(struct kvm_vcpu *vcpu)
          * The layout is as follows:
          * - gpr 2 contains the subchannel id (passed as addr)
          * - gpr 3 contains the virtqueue index (passed as datamatch)
+         * - gpr 4 contains the index on the bus (optionally)
          */
-        ret = kvm_io_bus_write(vcpu->kvm, KVM_VIRTIO_CCW_NOTIFY_BUS,
-                                vcpu->run->s.regs.gprs[2],
-                                8, &vcpu->run->s.regs.gprs[3]);
+        ret = kvm_io_bus_write_cookie(vcpu->kvm, KVM_VIRTIO_CCW_NOTIFY_BUS,
+                                      vcpu->run->s.regs.gprs[2],
+                                      8, &vcpu->run->s.regs.gprs[3],
+                                      vcpu->run->s.regs.gprs[4]);
         srcu_read_unlock(&vcpu->kvm->srcu, idx);
-        /* kvm_io_bus_write returns -EOPNOTSUPP if it found no match. */
+
+        /*
+         * Return cookie in gpr 2, but don't overwrite the register if the
+         * diagnose will be handled by userspace.
+         */
+        if (ret != -EOPNOTSUPP)
+                vcpu->run->s.regs.gprs[2] = ret;
+        /* kvm_io_bus_write_cookie returns -EOPNOTSUPP if it found no match. */
         return ret < 0 ? ret : 0;
 }
 
diff --git a/arch/s390/kvm/kvm-s390.c b/arch/s390/kvm/kvm-s390.c
index 34c1c9a90be2..776dafe918db 100644
--- a/arch/s390/kvm/kvm-s390.c
+++ b/arch/s390/kvm/kvm-s390.c
@@ -28,6 +28,7 @@
 #include <asm/pgtable.h>
 #include <asm/nmi.h>
 #include <asm/switch_to.h>
+#include <asm/facility.h>
 #include <asm/sclp.h>
 #include "kvm-s390.h"
 #include "gaccess.h"
@@ -84,9 +85,15 @@ struct kvm_stats_debugfs_item debugfs_entries[] = {
         { NULL }
 };
 
-static unsigned long long *facilities;
+unsigned long *vfacilities;
 static struct gmap_notifier gmap_notifier;
 
+/* test availability of vfacility */
+static inline int test_vfacility(unsigned long nr)
+{
+        return __test_facility(nr, (void *) vfacilities);
+}
+
 /* Section: not file related */
 int kvm_arch_hardware_enable(void *garbage)
 {
@@ -387,7 +394,7 @@ int kvm_arch_vcpu_setup(struct kvm_vcpu *vcpu)
         vcpu->arch.sie_block->ecb   = 6;
         vcpu->arch.sie_block->ecb2  = 8;
         vcpu->arch.sie_block->eca   = 0xC1002001U;
-        vcpu->arch.sie_block->fac   = (int) (long) facilities;
+        vcpu->arch.sie_block->fac   = (int) (long) vfacilities;
         hrtimer_init(&vcpu->arch.ckc_timer, CLOCK_REALTIME, HRTIMER_MODE_ABS);
         tasklet_init(&vcpu->arch.tasklet, kvm_s390_tasklet,
                      (unsigned long) vcpu);
@@ -1063,6 +1070,10 @@ int kvm_arch_create_memslot(struct kvm_memory_slot *slot, unsigned long npages)
         return 0;
 }
 
+void kvm_arch_memslots_updated(struct kvm *kvm)
+{
+}
+
 /* Section: memory related */
 int kvm_arch_prepare_memory_region(struct kvm *kvm,
                                    struct kvm_memory_slot *memslot,
@@ -1129,20 +1140,20 @@ static int __init kvm_s390_init(void)
          * to hold the maximum amount of facilities. On the other hand, we
          * only set facilities that are known to work in KVM.
          */
-        facilities = (unsigned long long *) get_zeroed_page(GFP_KERNEL|GFP_DMA);
-        if (!facilities) {
+        vfacilities = (unsigned long *) get_zeroed_page(GFP_KERNEL|GFP_DMA);
+        if (!vfacilities) {
                 kvm_exit();
                 return -ENOMEM;
         }
-        memcpy(facilities, S390_lowcore.stfle_fac_list, 16);
-        facilities[0] &= 0xff82fff3f47c0000ULL;
-        facilities[1] &= 0x001c000000000000ULL;
+        memcpy(vfacilities, S390_lowcore.stfle_fac_list, 16);
+        vfacilities[0] &= 0xff82fff3f47c0000UL;
+        vfacilities[1] &= 0x001c000000000000UL;
         return 0;
 }
 
 static void __exit kvm_s390_exit(void)
 {
-        free_page((unsigned long) facilities);
+        free_page((unsigned long) vfacilities);
         kvm_exit();
 }
 
diff --git a/arch/s390/kvm/kvm-s390.h b/arch/s390/kvm/kvm-s390.h
index 028ca9fd2158..dc99f1ca4267 100644
--- a/arch/s390/kvm/kvm-s390.h
+++ b/arch/s390/kvm/kvm-s390.h
@@ -24,6 +24,9 @@
 
 typedef int (*intercept_handler_t)(struct kvm_vcpu *vcpu);
 
+/* declare vfacilities extern */
+extern unsigned long *vfacilities;
+
 /* negativ values are error codes, positive values for internal conditions */
 #define SIE_INTERCEPT_RERUNVCPU         (1<<0)
 #define SIE_INTERCEPT_UCONTROL          (1<<1)
@@ -112,6 +115,13 @@ static inline u64 kvm_s390_get_base_disp_rs(struct kvm_vcpu *vcpu)
         return (base2 ? vcpu->run->s.regs.gprs[base2] : 0) + disp2;
 }
 
+/* Set the condition code in the guest program status word */
+static inline void kvm_s390_set_psw_cc(struct kvm_vcpu *vcpu, unsigned long cc)
+{
+        vcpu->arch.sie_block->gpsw.mask &= ~(3UL << 44);
+        vcpu->arch.sie_block->gpsw.mask |= cc << 44;
+}
+
 int kvm_s390_handle_wait(struct kvm_vcpu *vcpu);
 enum hrtimer_restart kvm_s390_idle_wakeup(struct hrtimer *timer);
 void kvm_s390_tasklet(unsigned long parm);
diff --git a/arch/s390/kvm/priv.c b/arch/s390/kvm/priv.c
index 4cdc54e63ebc..59200ee275e5 100644
--- a/arch/s390/kvm/priv.c
+++ b/arch/s390/kvm/priv.c
@@ -164,8 +164,7 @@ static int handle_tpi(struct kvm_vcpu *vcpu)
         kfree(inti);
 no_interrupt:
         /* Set condition code and we're done. */
-        vcpu->arch.sie_block->gpsw.mask &= ~(3ul << 44);
-        vcpu->arch.sie_block->gpsw.mask |= (cc & 3ul) << 44;
+        kvm_s390_set_psw_cc(vcpu, cc);
         return 0;
 }
 
@@ -220,15 +219,13 @@ static int handle_io_inst(struct kvm_vcpu *vcpu)
                  * Set condition code 3 to stop the guest from issueing channel
                  * I/O instructions.
                  */
-                vcpu->arch.sie_block->gpsw.mask &= ~(3ul << 44);
-                vcpu->arch.sie_block->gpsw.mask |= (3 & 3ul) << 44;
+                kvm_s390_set_psw_cc(vcpu, 3);
                 return 0;
         }
 }
 
 static int handle_stfl(struct kvm_vcpu *vcpu)
 {
-        unsigned int facility_list;
         int rc;
 
         vcpu->stat.instruction_stfl++;
@@ -236,15 +233,13 @@ static int handle_stfl(struct kvm_vcpu *vcpu)
         if (vcpu->arch.sie_block->gpsw.mask & PSW_MASK_PSTATE)
                 return kvm_s390_inject_program_int(vcpu, PGM_PRIVILEGED_OP);
 
-        /* only pass the facility bits, which we can handle */
-        facility_list = S390_lowcore.stfl_fac_list & 0xff82fff3;
-
         rc = copy_to_guest(vcpu, offsetof(struct _lowcore, stfl_fac_list),
-                           &facility_list, sizeof(facility_list));
+                           vfacilities, 4);
         if (rc)
                 return kvm_s390_inject_program_int(vcpu, PGM_ADDRESSING);
-        VCPU_EVENT(vcpu, 5, "store facility list value %x", facility_list);
-        trace_kvm_s390_handle_stfl(vcpu, facility_list);
+        VCPU_EVENT(vcpu, 5, "store facility list value %x",
+                   *(unsigned int *) vfacilities);
+        trace_kvm_s390_handle_stfl(vcpu, *(unsigned int *) vfacilities);
         return 0;
 }
 
@@ -387,7 +382,7 @@ static int handle_stsi(struct kvm_vcpu *vcpu)
                 return kvm_s390_inject_program_int(vcpu, PGM_PRIVILEGED_OP);
 
         if (fc > 3) {
-                vcpu->arch.sie_block->gpsw.mask |= 3ul << 44;     /* cc 3 */
+                kvm_s390_set_psw_cc(vcpu, 3);
                 return 0;
         }
 
@@ -397,7 +392,7 @@ static int handle_stsi(struct kvm_vcpu *vcpu)
 
         if (fc == 0) {
                 vcpu->run->s.regs.gprs[0] = 3 << 28;
-                vcpu->arch.sie_block->gpsw.mask &= ~(3ul << 44);  /* cc 0 */
+                kvm_s390_set_psw_cc(vcpu, 0);
                 return 0;
         }
 
@@ -431,12 +426,11 @@ static int handle_stsi(struct kvm_vcpu *vcpu)
         }
         trace_kvm_s390_handle_stsi(vcpu, fc, sel1, sel2, operand2);
         free_page(mem);
-        vcpu->arch.sie_block->gpsw.mask &= ~(3ul << 44);
+        kvm_s390_set_psw_cc(vcpu, 0);
         vcpu->run->s.regs.gprs[0] = 0;
         return 0;
 out_no_data:
-        /* condition code 3 */
-        vcpu->arch.sie_block->gpsw.mask |= 3ul << 44;
+        kvm_s390_set_psw_cc(vcpu, 3);
 out_exception:
         free_page(mem);
         return rc;
@@ -494,12 +488,12 @@ static int handle_epsw(struct kvm_vcpu *vcpu)
         kvm_s390_get_regs_rre(vcpu, &reg1, &reg2);
 
         /* This basically extracts the mask half of the psw. */
-        vcpu->run->s.regs.gprs[reg1] &= 0xffffffff00000000;
+        vcpu->run->s.regs.gprs[reg1] &= 0xffffffff00000000UL;
         vcpu->run->s.regs.gprs[reg1] |= vcpu->arch.sie_block->gpsw.mask >> 32;
         if (reg2) {
-                vcpu->run->s.regs.gprs[reg2] &= 0xffffffff00000000;
+                vcpu->run->s.regs.gprs[reg2] &= 0xffffffff00000000UL;
                 vcpu->run->s.regs.gprs[reg2] |=
-                        vcpu->arch.sie_block->gpsw.mask & 0x00000000ffffffff;
+                        vcpu->arch.sie_block->gpsw.mask & 0x00000000ffffffffUL;
         }
         return 0;
 }
diff --git a/arch/s390/mm/pgtable.c b/arch/s390/mm/pgtable.c
index 6d16132d0850..bf7c0dc64a76 100644
--- a/arch/s390/mm/pgtable.c
+++ b/arch/s390/mm/pgtable.c
@@ -335,7 +335,7 @@ int gmap_map_segment(struct gmap *gmap, unsigned long from,
 
         if ((from | to | len) & (PMD_SIZE - 1))
                 return -EINVAL;
-        if (len == 0 || from + len > PGDIR_SIZE ||
+        if (len == 0 || from + len > TASK_MAX_SIZE ||
             from + len < from || to + len < to)
                 return -EINVAL;
 
@@ -732,6 +732,11 @@ void gmap_do_ipte_notify(struct mm_struct *mm, unsigned long addr, pte_t *pte)
         spin_unlock(&gmap_notifier_lock);
 }
 
+static inline int page_table_with_pgste(struct page *page)
+{
+        return atomic_read(&page->_mapcount) == 0;
+}
+
 static inline unsigned long *page_table_alloc_pgste(struct mm_struct *mm,
                                                     unsigned long vmaddr)
 {
@@ -751,7 +756,7 @@ static inline unsigned long *page_table_alloc_pgste(struct mm_struct *mm,
         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_INVALID, PAGE_SIZE/2);
         clear_table(table + PTRS_PER_PTE, PGSTE_HR_BIT | PGSTE_HC_BIT,
@@ -818,6 +823,11 @@ EXPORT_SYMBOL(set_guest_storage_key);
 
 #else /* CONFIG_PGSTE */
 
+static inline int page_table_with_pgste(struct page *page)
+{
+        return 0;
+}
+
 static inline unsigned long *page_table_alloc_pgste(struct mm_struct *mm,
                                                     unsigned long vmaddr)
 {
@@ -894,12 +904,12 @@ void page_table_free(struct mm_struct *mm, unsigned long *table)
         struct page *page;
         unsigned int bit, mask;
 
-        if (mm_has_pgste(mm)) {
+        page = pfn_to_page(__pa(table) >> PAGE_SHIFT);
+        if (page_table_with_pgste(page)) {
                 gmap_disconnect_pgtable(mm, table);
                 return page_table_free_pgste(table);
         }
         /* Free 1K/2K page table fragment of a 4K page */
-        page = pfn_to_page(__pa(table) >> PAGE_SHIFT);
         bit = 1 << ((__pa(table) & ~PAGE_MASK)/(PTRS_PER_PTE*sizeof(pte_t)));
         spin_lock_bh(&mm->context.list_lock);
         if ((atomic_read(&page->_mapcount) & FRAG_MASK) != FRAG_MASK)
@@ -937,14 +947,14 @@ void page_table_free_rcu(struct mmu_gather *tlb, unsigned long *table)
         unsigned int bit, mask;
 
         mm = tlb->mm;
-        if (mm_has_pgste(mm)) {
+        page = pfn_to_page(__pa(table) >> PAGE_SHIFT);
+        if (page_table_with_pgste(page)) {
                 gmap_disconnect_pgtable(mm, table);
                 table = (unsigned long *) (__pa(table) | FRAG_MASK);
                 tlb_remove_table(tlb, table);
                 return;
         }
         bit = 1 << ((__pa(table) & ~PAGE_MASK) / (PTRS_PER_PTE*sizeof(pte_t)));
-        page = pfn_to_page(__pa(table) >> PAGE_SHIFT);
         spin_lock_bh(&mm->context.list_lock);
         if ((atomic_read(&page->_mapcount) & FRAG_MASK) != FRAG_MASK)
                 list_del(&page->lru);
@@ -1030,36 +1040,120 @@ void tlb_remove_table(struct mmu_gather *tlb, void *table)
 }
 
 #ifdef CONFIG_TRANSPARENT_HUGEPAGE
-void thp_split_vma(struct vm_area_struct *vma)
+static inline void thp_split_vma(struct vm_area_struct *vma)
 {
         unsigned long addr;
-        struct page *page;
 
-        for (addr = vma->vm_start; addr < vma->vm_end; addr += PAGE_SIZE) {
-                page = follow_page(vma, addr, FOLL_SPLIT);
-        }
+        for (addr = vma->vm_start; addr < vma->vm_end; addr += PAGE_SIZE)
+                follow_page(vma, addr, FOLL_SPLIT);
 }
 
-void thp_split_mm(struct mm_struct *mm)
+static inline void thp_split_mm(struct mm_struct *mm)
 {
-        struct vm_area_struct *vma = mm->mmap;
+        struct vm_area_struct *vma;
 
-        while (vma != NULL) {
+        for (vma = mm->mmap; vma != NULL; vma = vma->vm_next) {
                 thp_split_vma(vma);
                 vma->vm_flags &= ~VM_HUGEPAGE;
                 vma->vm_flags |= VM_NOHUGEPAGE;
-                vma = vma->vm_next;
         }
+        mm->def_flags |= VM_NOHUGEPAGE;
+}
+#else
+static inline void thp_split_mm(struct mm_struct *mm)
+{
 }
 #endif /* CONFIG_TRANSPARENT_HUGEPAGE */
 
+static unsigned long page_table_realloc_pmd(struct mmu_gather *tlb,
+                                struct mm_struct *mm, pud_t *pud,
+                                unsigned long addr, unsigned long end)
+{
+        unsigned long next, *table, *new;
+        struct page *page;
+        pmd_t *pmd;
+
+        pmd = pmd_offset(pud, addr);
+        do {
+                next = pmd_addr_end(addr, end);
+again:
+                if (pmd_none_or_clear_bad(pmd))
+                        continue;
+                table = (unsigned long *) pmd_deref(*pmd);
+                page = pfn_to_page(__pa(table) >> PAGE_SHIFT);
+                if (page_table_with_pgste(page))
+                        continue;
+                /* Allocate new page table with pgstes */
+                new = page_table_alloc_pgste(mm, addr);
+                if (!new) {
+                        mm->context.has_pgste = 0;
+                        continue;
+                }
+                spin_lock(&mm->page_table_lock);
+                if (likely((unsigned long *) pmd_deref(*pmd) == table)) {
+                        /* Nuke pmd entry pointing to the "short" page table */
+                        pmdp_flush_lazy(mm, addr, pmd);
+                        pmd_clear(pmd);
+                        /* Copy ptes from old table to new table */
+                        memcpy(new, table, PAGE_SIZE/2);
+                        clear_table(table, _PAGE_INVALID, PAGE_SIZE/2);
+                        /* Establish new table */
+                        pmd_populate(mm, pmd, (pte_t *) new);
+                        /* Free old table with rcu, there might be a walker! */
+                        page_table_free_rcu(tlb, table);
+                        new = NULL;
+                }
+                spin_unlock(&mm->page_table_lock);
+                if (new) {
+                        page_table_free_pgste(new);
+                        goto again;
+                }
+        } while (pmd++, addr = next, addr != end);
+
+        return addr;
+}
+
+static unsigned long page_table_realloc_pud(struct mmu_gather *tlb,
+                                   struct mm_struct *mm, pgd_t *pgd,
+                                   unsigned long addr, unsigned long end)
+{
+        unsigned long next;
+        pud_t *pud;
+
+        pud = pud_offset(pgd, addr);
+        do {
+                next = pud_addr_end(addr, end);
+                if (pud_none_or_clear_bad(pud))
+                        continue;
+                next = page_table_realloc_pmd(tlb, mm, pud, addr, next);
+        } while (pud++, addr = next, addr != end);
+
+        return addr;
+}
+
+static void page_table_realloc(struct mmu_gather *tlb, struct mm_struct *mm,
+                               unsigned long addr, unsigned long end)
+{
+        unsigned long next;
+        pgd_t *pgd;
+
+        pgd = pgd_offset(mm, addr);
+        do {
+                next = pgd_addr_end(addr, end);
+                if (pgd_none_or_clear_bad(pgd))
+                        continue;
+                next = page_table_realloc_pud(tlb, mm, pgd, addr, next);
+        } while (pgd++, addr = next, addr != end);
+}
+
 /*
  * switch on pgstes for its userspace process (for kvm)
  */
 int s390_enable_sie(void)
 {
         struct task_struct *tsk = current;
-        struct mm_struct *mm, *old_mm;
+        struct mm_struct *mm = tsk->mm;
+        struct mmu_gather tlb;
 
         /* Do we have switched amode? If no, we cannot do sie */
         if (s390_user_mode == HOME_SPACE_MODE)
@@ -1069,57 +1163,16 @@ int s390_enable_sie(void)
         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, TASK_SIZE);
+        page_table_realloc(&tlb, mm, 0, TASK_SIZE);
+        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);