1 files changed, 13 insertions, 8 deletions

diff --git a/arch/s390/mm/pgtable.c b/arch/s390/mm/pgtable.c
index 94f37a9fb1e5..a9be08899b0c 100644
--- a/arch/s390/mm/pgtable.c
+++ b/arch/s390/mm/pgtable.c
@@ -1087,10 +1087,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 +1127,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 +1146,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;
 }
 
 /*
@@ -1165,9 +1170,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;