diff --git a/arch/arm64/include/asm/pgtable.h b/arch/arm64/include/asm/pgtable.h
index f914c30b74871974176a9239ab4b4c23b09f1527..91e2c5a1bf29a1ae645048129a00c963fd76294b 100644
--- a/arch/arm64/include/asm/pgtable.h
+++ b/arch/arm64/include/asm/pgtable.h
@@ -1043,6 +1043,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/sparse-vmemmap.c b/mm/sparse-vmemmap.c
index a47b027af1f73c8d29912d702ff98d4d0099279a..b2536a3e5e2d23b0a5517fc6bf4d8c97c191af26 100644
--- a/mm/sparse-vmemmap.c
+++ b/mm/sparse-vmemmap.c
@@ -70,6 +70,14 @@ static inline void vmemmap_update_pte(unsigned long addr,
 }
 #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
+
 #ifndef vmemmap_flush_tlb_all
 static inline void vmemmap_flush_tlb_all(void)
 {
@@ -107,7 +115,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
@@ -121,10 +129,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;
 }