diff --git a/arch/x86/kernel/process_64.c b/arch/x86/kernel/process_64.c
index e9f7cfdb94202..226472332a70d 100644
--- a/arch/x86/kernel/process_64.c
+++ b/arch/x86/kernel/process_64.c
@@ -812,6 +812,9 @@ static void enable_lam_func(void *__mm)
 
 static void mm_enable_lam(struct mm_struct *mm)
 {
+	mm->context.lam_cr3_mask = X86_CR3_LAM_U57;
+	mm->context.untag_mask =  ~GENMASK(62, 57);
+
 	/*
 	 * Even though the process must still be single-threaded at this
 	 * point, kernel threads may be using the mm.  IPI those kernel
@@ -846,13 +849,7 @@ static int prctl_enable_tagged_addr(struct mm_struct *mm, unsigned long nr_bits)
 		return -EBUSY;
 	}
 
-	if (!nr_bits) {
-		mmap_write_unlock(mm);
-		return -EINVAL;
-	} else if (nr_bits <= LAM_U57_BITS) {
-		mm->context.lam_cr3_mask = X86_CR3_LAM_U57;
-		mm->context.untag_mask =  ~GENMASK(62, 57);
-	} else {
+	if (!nr_bits || nr_bits > LAM_U57_BITS) {
 		mmap_write_unlock(mm);
 		return -EINVAL;
 	}