1 files changed, 31 insertions, 12 deletions

diff --git a/arch/s390/mm/pgtable.c b/arch/s390/mm/pgtable.c
index de8cbc30dcd..0a2e5e08674 100644
--- a/arch/s390/mm/pgtable.c
+++ b/arch/s390/mm/pgtable.c
@@ -48,12 +48,23 @@ void crst_table_free(struct mm_struct *mm, unsigned long *table)
 }
 
 #ifdef CONFIG_64BIT
+static void __crst_table_upgrade(void *arg)
+{
+	struct mm_struct *mm = arg;
+
+	if (current->active_mm == mm)
+		update_mm(mm, current);
+	__tlb_flush_local();
+}
+
 int crst_table_upgrade(struct mm_struct *mm, unsigned long limit)
 {
 	unsigned long *table, *pgd;
 	unsigned long entry;
+	int flush;
 
 	BUG_ON(limit > (1UL << 53));
+	flush = 0;
 repeat:
 	table = crst_table_alloc(mm);
 	if (!table)
@@ -79,12 +90,15 @@ repeat:
 		mm->pgd = (pgd_t *) table;
 		mm->task_size = mm->context.asce_limit;
 		table = NULL;
+		flush = 1;
 	}
 	spin_unlock_bh(&mm->page_table_lock);
 	if (table)
 		crst_table_free(mm, table);
 	if (mm->context.asce_limit < limit)
 		goto repeat;
+	if (flush)
+		on_each_cpu(__crst_table_upgrade, mm, 0);
 	return 0;
 }
 
@@ -92,6 +106,8 @@ void crst_table_downgrade(struct mm_struct *mm, unsigned long limit)
 {
 	pgd_t *pgd;
 
+	if (current->active_mm == mm)
+		__tlb_flush_mm(mm);
 	while (mm->context.asce_limit > limit) {
 		pgd = mm->pgd;
 		switch (pgd_val(*pgd) & _REGION_ENTRY_TYPE_MASK) {
@@ -114,6 +130,8 @@ void crst_table_downgrade(struct mm_struct *mm, unsigned long limit)
 		mm->task_size = mm->context.asce_limit;
 		crst_table_free(mm, (unsigned long *) pgd);
 	}
+	if (current->active_mm == mm)
+		update_mm(mm, current);
 }
 #endif
 
@@ -1087,10 +1105,9 @@ again:
 			continue;
 		/* Allocate new page table with pgstes */
 		new = page_table_alloc_pgste(mm, addr);
-		if (!new) {
-			mm->context.has_pgste = 0;
-			continue;
-		}
+		if (!new)
+			return -ENOMEM;
+
 		spin_lock(&mm->page_table_lock);
 		if (likely((unsigned long *) pmd_deref(*pmd) == table)) {
 			/* Nuke pmd entry pointing to the "short" page table */
@@ -1128,13 +1145,15 @@ static unsigned long page_table_realloc_pud(struct mmu_gather *tlb,
 		if (pud_none_or_clear_bad(pud))
 			continue;
 		next = page_table_realloc_pmd(tlb, mm, pud, addr, next);
+		if (unlikely(IS_ERR_VALUE(next)))
+			return 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)
+static unsigned long page_table_realloc(struct mmu_gather *tlb, struct mm_struct *mm,
+					unsigned long addr, unsigned long end)
 {
 	unsigned long next;
 	pgd_t *pgd;
@@ -1145,7 +1164,11 @@ static void page_table_realloc(struct mmu_gather *tlb, struct mm_struct *mm,
 		if (pgd_none_or_clear_bad(pgd))
 			continue;
 		next = page_table_realloc_pud(tlb, mm, pgd, addr, next);
+		if (unlikely(IS_ERR_VALUE(next)))
+			return next;
 	} while (pgd++, addr = next, addr != end);
+
+	return 0;
 }
 
 /*
@@ -1157,10 +1180,6 @@ int s390_enable_sie(void)
 	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)
-		return -EINVAL;
-
 	/* Do we have pgstes? if yes, we are done */
 	if (mm_has_pgste(tsk->mm))
 		return 0;
@@ -1169,9 +1188,9 @@ int s390_enable_sie(void)
 	/* split thp mappings and disable thp for future mappings */
 	thp_split_mm(mm);
 	/* 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);
+	if (!page_table_realloc(&tlb, mm, 0, TASK_SIZE))
+		mm->context.has_pgste = 1;
 	tlb_finish_mmu(&tlb, 0, TASK_SIZE);
 	up_write(&mm->mmap_sem);
 	return mm->context.has_pgste ? 0 : -ENOMEM;