{
        unsigned long ps = memparse(opt, &opt);
 
-       if (ps == PMD_SIZE) {
-               hugetlb_add_hstate(PMD_SHIFT - PAGE_SHIFT);
-       } else if (ps == PUD_SIZE) {
-               hugetlb_add_hstate(PUD_SHIFT - PAGE_SHIFT);
-       } else if (ps == (PAGE_SIZE * CONT_PTES)) {
-               hugetlb_add_hstate(CONT_PTE_SHIFT);
-       } else if (ps == (PMD_SIZE * CONT_PMDS)) {
-               hugetlb_add_hstate((PMD_SHIFT + CONT_PMD_SHIFT) - PAGE_SHIFT);
-       } else {
-               hugetlb_bad_size();
-               pr_err("hugepagesz: Unsupported page size %lu K\n", ps >> 10);
-               return 0;
+       switch (ps) {
+#ifdef CONFIG_ARM64_4K_PAGES
+       case PUD_SIZE:
+#endif
+       case PMD_SIZE * CONT_PMDS:
+       case PMD_SIZE:
+       case PAGE_SIZE * CONT_PTES:
+               hugetlb_add_hstate(ilog2(ps) - PAGE_SHIFT);
+               return 1;
        }
-       return 1;
+
+       hugetlb_bad_size();
+       pr_err("hugepagesz: Unsupported page size %lu K\n", ps >> 10);
+       return 0;
 }
 __setup("hugepagesz=", setup_hugepagesz);