diff --git a/fs/userfaultfd.c b/fs/userfaultfd.c
index d27b66eb5aa0873a0e03e62eb944037a2811ec5d..1222e507fade7f6cd8d8d923bc06676235681591 100644
--- a/fs/userfaultfd.c
+++ b/fs/userfaultfd.c
@@ -335,8 +335,7 @@ static inline bool userfaultfd_must_wait(struct userfaultfd_ctx *ctx,
 	 * changes under us.
 	 */
 #ifdef CONFIG_USERSWAP
-	if ((reason & VM_USWAP) && (!pte_present(*pte)))
-		ret = true;
+	uswap_must_wait(reason, *pte, &ret);
 #endif
 	if (pte_none(*pte))
 		ret = true;
@@ -875,8 +874,7 @@ static int userfaultfd_release(struct inode *inode, struct file *file)
 	for (vma = mm->mmap; vma; vma = vma->vm_next) {
 		userfault_flags = VM_UFFD_MISSING | VM_UFFD_WP;
 #ifdef CONFIG_USERSWAP
-		if (enable_userswap)
-			userfault_flags |= VM_USWAP;
+		uswap_release(&userfault_flags);
 #endif
 		cond_resched();
 		BUG_ON(!!vma->vm_userfaultfd_ctx.ctx ^
@@ -1297,26 +1295,8 @@ static int userfaultfd_register(struct userfaultfd_ctx *ctx,
 		goto out;
 	vm_flags = 0;
 #ifdef CONFIG_USERSWAP
-	/*
-	 * register the whole vma overlapping with the address range to avoid
-	 * splitting the vma.
-	 */
-	if (enable_userswap && (uffdio_register.mode & UFFDIO_REGISTER_MODE_USWAP)) {
-		uffdio_register.mode &= ~UFFDIO_REGISTER_MODE_USWAP;
-		if (!uffdio_register.mode)
-			goto out;
-		vm_flags |= VM_USWAP;
-		end = uffdio_register.range.start + uffdio_register.range.len - 1;
-		vma = find_vma(mm, uffdio_register.range.start);
-		if (!vma)
-			goto out;
-		uffdio_register.range.start = vma->vm_start;
-
-		vma = find_vma(mm, end);
-		if (!vma)
-			goto out;
-		uffdio_register.range.len = vma->vm_end - uffdio_register.range.start;
-	}
+	if (!uswap_register(&uffdio_register, &vm_flags, mm))
+		goto out;
 #endif
 	if (uffdio_register.mode & ~(UFFDIO_REGISTER_MODE_MISSING|
 				     UFFDIO_REGISTER_MODE_WP))
@@ -2041,15 +2021,6 @@ SYSCALL_DEFINE1(userfaultfd, int, flags)
 	return fd;
 }
 
-#ifdef CONFIG_USERSWAP
-static int __init enable_userswap_setup(char *str)
-{
-	enable_userswap = true;
-	return 1;
-}
-__setup("enable_userswap", enable_userswap_setup);
-#endif
-
 static int __init userfaultfd_init(void)
 {
 	userfaultfd_ctx_cachep = kmem_cache_create("userfaultfd_ctx_cache",
diff --git a/include/linux/userswap.h b/include/linux/userswap.h
index 9d1d0ee2e9d757346aaa1fe52ef71f25b6f43a9e..14826e3d9c2c60fa05bc70c96bbfdf37dd3ca91e 100644
--- a/include/linux/userswap.h
+++ b/include/linux/userswap.h
@@ -28,6 +28,12 @@ int mfill_atomic_pte_nocopy(struct mm_struct *dst_mm,
 unsigned long uswap_mremap(unsigned long old_addr, unsigned long old_len,
 			   unsigned long new_addr, unsigned long new_len);
 
+bool uswap_register(struct uffdio_register *uffdio_register,
+		    unsigned long *vm_flags, struct mm_struct *mm);
+
+bool do_uswap_page(swp_entry_t entry, struct vm_fault *vmf,
+		   struct vm_area_struct *vma, vm_fault_t *ret);
+
 static inline bool uswap_check_copy_mode(struct vm_area_struct *vma, __u64 mode)
 {
 	if (!(vma->vm_flags & VM_USWAP) && (mode & UFFDIO_COPY_MODE_DIRECT_MAP))
@@ -72,6 +78,18 @@ static inline void uswap_get_cpu_id(unsigned long reason, struct uffd_msg *msg)
 		msg->reserved3 = smp_processor_id();
 }
 
+static inline void uswap_release(unsigned long *userfault_flags)
+{
+	if (enable_userswap)
+		*userfault_flags |= VM_USWAP;
+}
+
+static inline void uswap_must_wait(unsigned long reason, pte_t pte, bool *ret)
+{
+	if ((reason & VM_USWAP) && (!pte_present(pte)))
+		*ret = true;
+}
+
 #endif /* CONFIG_USERSWAP */
 
 #endif /* _LINUX_USERSWAP_H */
diff --git a/mm/memory.c b/mm/memory.c
index ed017586db1f5269aae54bca738ac925366f6497..5941a4f4ea4b154243a368f73f2925814f1cf45d 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -3396,22 +3396,8 @@ vm_fault_t do_swap_page(struct vm_fault *vmf)
 
 	entry = pte_to_swp_entry(vmf->orig_pte);
 #ifdef CONFIG_USERSWAP
-	if (swp_type(entry) == SWP_USERSWAP_ENTRY) {
-		/* print error if we come across a nested fault */
-		if (!strncmp(current->comm, "uswap", 5)) {
-			pr_err("USWAP: fault %lx is triggered by %s\n",
-					vmf->address, current->comm);
-			return VM_FAULT_SIGBUS;
-		}
-		if (!(vma->vm_flags & VM_UFFD_MISSING)) {
-			pr_err("USWAP: addr %lx flags %lx is not a user swap page",
-					vmf->address, vma->vm_flags);
-			goto skip_uswap;
-		}
-		ret = handle_userfault(vmf, VM_UFFD_MISSING | VM_USWAP);
+	if (!do_uswap_page(entry, vmf, vma, &ret))
 		return ret;
-	}
-skip_uswap:
 #endif
 	if (unlikely(non_swap_entry(entry))) {
 		if (is_migration_entry(entry)) {
diff --git a/mm/userswap.c b/mm/userswap.c
index dd212b1a02e6ed72b1d1864c1b9c833447690f9e..384548db9ba2cf9da51d53cb57a00e9aa490e8b8 100644
--- a/mm/userswap.c
+++ b/mm/userswap.c
@@ -10,6 +10,7 @@
 #include <linux/rmap.h>
 #include <linux/mmu_notifier.h>
 #include <linux/userswap.h>
+#include <linux/userfaultfd_k.h>
 
 #include "internal.h"
 
@@ -458,3 +459,67 @@ int mfill_atomic_pte_nocopy(struct mm_struct *mm,
 	put_page(page);
 	return ret;
 }
+
+/*
+ * register the whole vma overlapping with the address range to avoid splitting
+ * the vma.
+ */
+bool uswap_register(struct uffdio_register *uffdio_register,
+		    unsigned long *vm_flags, struct mm_struct *mm)
+{
+	struct vm_area_struct *vma;
+	unsigned long end;
+
+	if (!enable_userswap)
+		return true;
+	if (!(uffdio_register->mode & UFFDIO_REGISTER_MODE_USWAP))
+		return true;
+	uffdio_register->mode &= ~UFFDIO_REGISTER_MODE_USWAP;
+	if (!uffdio_register->mode)
+		return false;
+
+	end = uffdio_register->range.start + uffdio_register->range.len - 1;
+	vma = find_vma(mm, uffdio_register->range.start);
+	if (!vma)
+		return false;
+	uffdio_register->range.start = vma->vm_start;
+	vma = find_vma(mm, end);
+	if (!vma)
+		return false;
+	uffdio_register->range.len = vma->vm_end - uffdio_register->range.start;
+
+	*vm_flags |= VM_USWAP;
+
+	return true;
+}
+
+bool do_uswap_page(swp_entry_t entry, struct vm_fault *vmf,
+		   struct vm_area_struct *vma, vm_fault_t *ret)
+{
+	if (swp_type(entry) != SWP_USERSWAP_ENTRY)
+		return true;
+
+	/* print error if we come across a nested fault */
+	if (!strncmp(current->comm, "uswap", 5)) {
+		pr_err("USWAP: fault %lx is triggered by %s\n", vmf->address,
+		       current->comm);
+		*ret = VM_FAULT_SIGBUS;
+		return false;
+	}
+
+	if (!(vma->vm_flags & VM_UFFD_MISSING)) {
+		pr_err("USWAP: addr %lx flags %lx is not a user swap page",
+				vmf->address, vma->vm_flags);
+		return true;
+	}
+
+	*ret = handle_userfault(vmf, VM_UFFD_MISSING | VM_USWAP);
+	return false;
+}
+
+static int __init enable_userswap_setup(char *str)
+{
+	enable_userswap = true;
+	return 1;
+}
+__setup("enable_userswap", enable_userswap_setup);