diff --git a/arch/arm64/include/asm/pgtable.h b/arch/arm64/include/asm/pgtable.h
index 626e43967e0a310a305739c900bce2813e48b8eb..2b66aab73dbc42b7eda8c8833beccd6410c32d8e 100644
--- a/arch/arm64/include/asm/pgtable.h
+++ b/arch/arm64/include/asm/pgtable.h
@@ -1515,6 +1515,8 @@ void vmemmap_update_pmd(unsigned long addr, pmd_t *pmdp, pte_t *ptep);
 #define vmemmap_update_pmd vmemmap_update_pmd
 void vmemmap_update_pte(unsigned long addr, pte_t *ptep, pte_t pte);
 #define vmemmap_update_pte vmemmap_update_pte
+#define vmemmap_split_lock(lock)	spin_lock_irq(lock)
+#define vmemmap_split_unlock(lock)	spin_unlock_irq(lock)
 #endif
 
 #endif /* !__ASSEMBLY__ */
diff --git a/mm/hugetlb_vmemmap.c b/mm/hugetlb_vmemmap.c
index 149ab629855c11ebec2da97d1ce3c998820ea708..2bde429b2ea3dd11d5a47ad266a4db9eabf08f18 100644
--- a/mm/hugetlb_vmemmap.c
+++ b/mm/hugetlb_vmemmap.c
@@ -69,6 +69,14 @@ static inline void vmemmap_flush_tlb_range(unsigned long start,
 }
 #endif
 
+#ifndef vmemmap_split_lock
+#define vmemmap_split_lock(lock)	spin_lock(lock)
+#endif
+
+#ifndef vmemmap_split_unlock
+#define vmemmap_split_unlock(lock)	spin_unlock(lock)
+#endif
+
 static int split_vmemmap_huge_pmd(pmd_t *pmd, unsigned long start)
 {
 	pmd_t __pmd;
@@ -99,7 +107,7 @@ static int split_vmemmap_huge_pmd(pmd_t *pmd, unsigned long start)
 		set_pte_at(&init_mm, addr, pte, entry);
 	}
 
-	spin_lock_irq(&init_mm.page_table_lock);
+	vmemmap_split_lock(&init_mm.page_table_lock);
 	if (likely(pmd_leaf(*pmd))) {
 		/*
 		 * Higher order allocations from buddy allocator must be able to
@@ -113,10 +121,12 @@ static int split_vmemmap_huge_pmd(pmd_t *pmd, unsigned long start)
 		smp_wmb();
 		vmemmap_update_pmd(start, pmd, pgtable);
 		vmemmap_flush_tlb_range(start, start + PMD_SIZE);
-	} else {
-		pte_free_kernel(&init_mm, pgtable);
+		pgtable = NULL;
 	}
-	spin_unlock_irq(&init_mm.page_table_lock);
+	vmemmap_split_unlock(&init_mm.page_table_lock);
+
+	if (unlikely(pgtable))
+		pte_free_kernel(&init_mm, pgtable);
 
 	return 0;
 }