#include "internal.h"
 #include "pgalloc-track.h"
 
+#ifdef CONFIG_HAVE_ARCH_HUGE_VMALLOC
+static bool __ro_after_init vmap_allow_huge = true;
+
+static int __init set_nohugevmalloc(char *str)
+{
+       vmap_allow_huge = false;
+       return 0;
+}
+early_param("nohugevmalloc", set_nohugevmalloc);
+#else /* CONFIG_HAVE_ARCH_HUGE_VMALLOC */
+static const bool vmap_allow_huge = false;
+#endif /* CONFIG_HAVE_ARCH_HUGE_VMALLOC */
+
 bool is_vmalloc_addr(const void *x)
 {
        unsigned long addr = (unsigned long)x;
        return 0;
 }
 
-/**
- * map_kernel_range_noflush - map kernel VM area with the specified pages
- * @addr: start of the VM area to map
- * @size: size of the VM area to map
- * @prot: page protection flags to use
- * @pages: pages to map
- *
- * Map PFN_UP(@size) pages at @addr.  The VM area @addr and @size specify should
- * have been allocated using get_vm_area() and its friends.
- *
- * NOTE:
- * This function does NOT do any cache flushing.  The caller is responsible for
- * calling flush_cache_vmap() on to-be-mapped areas before calling this
- * function.
- *
- * RETURNS:
- * 0 on success, -errno on failure.
- */
-int map_kernel_range_noflush(unsigned long addr, unsigned long size,
-                            pgprot_t prot, struct page **pages)
+static int vmap_small_pages_range_noflush(unsigned long addr, unsigned long end,
+               pgprot_t prot, struct page **pages)
 {
        unsigned long start = addr;
-       unsigned long end = addr + size;
-       unsigned long next;
        pgd_t *pgd;
+       unsigned long next;
        int err = 0;
        int nr = 0;
        pgtbl_mod_mask mask = 0;
        return 0;
 }
 
+static int vmap_pages_range_noflush(unsigned long addr, unsigned long end,
+               pgprot_t prot, struct page **pages, unsigned int page_shift)
+{
+       unsigned int i, nr = (end - addr) >> PAGE_SHIFT;
+
+       WARN_ON(page_shift < PAGE_SHIFT);
+
+       if (!IS_ENABLED(CONFIG_HAVE_ARCH_HUGE_VMALLOC) ||
+                       page_shift == PAGE_SHIFT)
+               return vmap_small_pages_range_noflush(addr, end, prot, pages);
+
+       for (i = 0; i < nr; i += 1U << (page_shift - PAGE_SHIFT)) {
+               int err;
+
+               err = vmap_range_noflush(addr, addr + (1UL << page_shift),
+                                       __pa(page_address(pages[i])), prot,
+                                       page_shift);
+               if (err)
+                       return err;
+
+               addr += 1UL << page_shift;
+       }
+
+       return 0;
+}
+
+static int vmap_pages_range(unsigned long addr, unsigned long end,
+               pgprot_t prot, struct page **pages, unsigned int page_shift)
+{
+       int err;
+
+       err = vmap_pages_range_noflush(addr, end, prot, pages, page_shift);
+       flush_cache_vmap(addr, end);
+       return err;
+}
+
+/**
+ * map_kernel_range_noflush - map kernel VM area with the specified pages
+ * @addr: start of the VM area to map
+ * @size: size of the VM area to map
+ * @prot: page protection flags to use
+ * @pages: pages to map
+ *
+ * Map PFN_UP(@size) pages at @addr.  The VM area @addr and @size specify should
+ * have been allocated using get_vm_area() and its friends.
+ *
+ * NOTE:
+ * This function does NOT do any cache flushing.  The caller is responsible for
+ * calling flush_cache_vmap() on to-be-mapped areas before calling this
+ * function.
+ *
+ * RETURNS:
+ * 0 on success, -errno on failure.
+ */
+int map_kernel_range_noflush(unsigned long addr, unsigned long size,
+                            pgprot_t prot, struct page **pages)
+{
+       return vmap_pages_range_noflush(addr, addr + size, prot, pages, PAGE_SHIFT);
+}
+
 int map_kernel_range(unsigned long start, unsigned long size, pgprot_t prot,
                struct page **pages)
 {
 
 static struct vm_struct *vmlist __initdata;
 
+static inline unsigned int vm_area_page_order(struct vm_struct *vm)
+{
+#ifdef CONFIG_HAVE_ARCH_HUGE_VMALLOC
+       return vm->page_order;
+#else
+       return 0;
+#endif
+}
+
+static inline void set_vm_area_page_order(struct vm_struct *vm, unsigned int order)
+{
+#ifdef CONFIG_HAVE_ARCH_HUGE_VMALLOC
+       vm->page_order = order;
+#else
+       BUG_ON(order != 0);
+#endif
+}
+
 /**
  * vm_area_add_early - add vmap area early during boot
  * @vm: vm_struct to add
 {
        int i;
 
+       /* HUGE_VMALLOC passes small pages to set_direct_map */
        for (i = 0; i < area->nr_pages; i++)
                if (page_address(area->pages[i]))
                        set_direct_map(area->pages[i]);
 static void vm_remove_mappings(struct vm_struct *area, int deallocate_pages)
 {
        unsigned long start = ULONG_MAX, end = 0;
+       unsigned int page_order = vm_area_page_order(area);
        int flush_reset = area->flags & VM_FLUSH_RESET_PERMS;
        int flush_dmap = 0;
        int i;
         * map. Find the start and end range of the direct mappings to make sure
         * the vm_unmap_aliases() flush includes the direct map.
         */
-       for (i = 0; i < area->nr_pages; i++) {
+       for (i = 0; i < area->nr_pages; i += 1U << page_order) {
                unsigned long addr = (unsigned long)page_address(area->pages[i]);
                if (addr) {
+                       unsigned long page_size;
+
+                       page_size = PAGE_SIZE << page_order;
                        start = min(addr, start);
-                       end = max(addr + PAGE_SIZE, end);
+                       end = max(addr + page_size, end);
                        flush_dmap = 1;
                }
        }
        vm_remove_mappings(area, deallocate_pages);
 
        if (deallocate_pages) {
+               unsigned int page_order = vm_area_page_order(area);
                int i;
 
-               for (i = 0; i < area->nr_pages; i++) {
+               for (i = 0; i < area->nr_pages; i += 1U << page_order) {
                        struct page *page = area->pages[i];
 
                        BUG_ON(!page);
-                       __free_pages(page, 0);
+                       __free_pages(page, page_order);
                }
                atomic_long_sub(area->nr_pages, &nr_vmalloc_pages);
 
 #endif /* CONFIG_VMAP_PFN */
 
 static void *__vmalloc_area_node(struct vm_struct *area, gfp_t gfp_mask,
-                                pgprot_t prot, int node)
+                                pgprot_t prot, unsigned int page_shift,
+                                int node)
 {
        const gfp_t nested_gfp = (gfp_mask & GFP_RECLAIM_MASK) | __GFP_ZERO;
-       unsigned int nr_pages = get_vm_area_size(area) >> PAGE_SHIFT;
+       unsigned long addr = (unsigned long)area->addr;
+       unsigned long size = get_vm_area_size(area);
        unsigned long array_size;
-       unsigned int i;
+       unsigned int nr_small_pages = size >> PAGE_SHIFT;
+       unsigned int page_order;
        struct page **pages;
+       unsigned int i;
 
-       array_size = (unsigned long)nr_pages * sizeof(struct page *);
+       array_size = (unsigned long)nr_small_pages * sizeof(struct page *);
        gfp_mask |= __GFP_NOWARN;
        if (!(gfp_mask & (GFP_DMA | GFP_DMA32)))
                gfp_mask |= __GFP_HIGHMEM;
        }
 
        area->pages = pages;
-       area->nr_pages = nr_pages;
+       area->nr_pages = nr_small_pages;
+       set_vm_area_page_order(area, page_shift - PAGE_SHIFT);
 
-       for (i = 0; i < area->nr_pages; i++) {
-               struct page *page;
+       page_order = vm_area_page_order(area);
 
-               if (node == NUMA_NO_NODE)
-                       page = alloc_page(gfp_mask);
-               else
-                       page = alloc_pages_node(node, gfp_mask, 0);
+       /*
+        * Careful, we allocate and map page_order pages, but tracking is done
+        * per PAGE_SIZE page so as to keep the vm_struct APIs independent of
+        * the physical/mapped size.
+        */
+       for (i = 0; i < area->nr_pages; i += 1U << page_order) {
+               struct page *page;
+               int p;
 
+               /* Compound pages required for remap_vmalloc_page */
+               page = alloc_pages_node(node, gfp_mask | __GFP_COMP, page_order);
                if (unlikely(!page)) {
                        /* Successfully allocated i pages, free them in __vfree() */
                        area->nr_pages = i;
                        atomic_long_add(area->nr_pages, &nr_vmalloc_pages);
                        goto fail;
                }
-               area->pages[i] = page;
+
+               for (p = 0; p < (1U << page_order); p++)
+                       area->pages[i + p] = page + p;
+
                if (gfpflags_allow_blocking(gfp_mask))
                        cond_resched();
        }
        atomic_long_add(area->nr_pages, &nr_vmalloc_pages);
 
-       if (map_kernel_range((unsigned long)area->addr, get_vm_area_size(area),
-                       prot, pages) < 0)
+       if (vmap_pages_range(addr, addr + size, prot, pages, page_shift) < 0)
                goto fail;
 
        return area->addr;
 fail:
        warn_alloc(gfp_mask, NULL,
                          "vmalloc: allocation failure, allocated %ld of %ld bytes",
-                         (area->nr_pages*PAGE_SIZE), area->size);
+                         (area->nr_pages*PAGE_SIZE), size);
        __vfree(area->addr);
        return NULL;
 }
        struct vm_struct *area;
        void *addr;
        unsigned long real_size = size;
+       unsigned long real_align = align;
+       unsigned int shift = PAGE_SHIFT;
 
-       size = PAGE_ALIGN(size);
-       if (!size || (size >> PAGE_SHIFT) > totalram_pages())
+       if (!size || (size >> PAGE_SHIFT) > totalram_pages()) {
+               area = NULL;
                goto fail;
+       }
+
+       if (vmap_allow_huge && !(vm_flags & VM_NO_HUGE_VMAP) &&
+                       arch_vmap_pmd_supported(prot)) {
+               unsigned long size_per_node;
 
-       area = __get_vm_area_node(real_size, align, VM_ALLOC | VM_UNINITIALIZED |
+               /*
+                * Try huge pages. Only try for PAGE_KERNEL allocations,
+                * others like modules don't yet expect huge pages in
+                * their allocations due to apply_to_page_range not
+                * supporting them.
+                */
+
+               size_per_node = size;
+               if (node == NUMA_NO_NODE)
+                       size_per_node /= num_online_nodes();
+               if (size_per_node >= PMD_SIZE) {
+                       shift = PMD_SHIFT;
+                       align = max(real_align, 1UL << shift);
+                       size = ALIGN(real_size, 1UL << shift);
+               }
+       }
+
+again:
+       size = PAGE_ALIGN(size);
+       area = __get_vm_area_node(size, align, VM_ALLOC | VM_UNINITIALIZED |
                                vm_flags, start, end, node, gfp_mask, caller);
        if (!area)
                goto fail;
 
-       addr = __vmalloc_area_node(area, gfp_mask, prot, node);
+       addr = __vmalloc_area_node(area, gfp_mask, prot, shift, node);
        if (!addr)
-               return NULL;
+               goto fail;
 
        /*
         * In this function, newly allocated vm_struct has VM_UNINITIALIZED
        return addr;
 
 fail:
-       warn_alloc(gfp_mask, NULL,
+       if (shift > PAGE_SHIFT) {
+               shift = PAGE_SHIFT;
+               align = real_align;
+               size = real_size;
+               goto again;
+       }
+
+       if (!area) {
+               /* Warn for area allocation, page allocations already warn */
+               warn_alloc(gfp_mask, NULL,
                          "vmalloc: allocation failure: %lu bytes", real_size);
+       }
        return NULL;
 }