#include <linux/extable.h>
 #include <linux/log2.h>
 #include <linux/bpf_verifier.h>
+#include <linux/nodemask.h>
 
 #include <asm/barrier.h>
 #include <asm/unaligned.h>
  * allocator. The prog_pack allocator uses HPAGE_PMD_SIZE page (2MB on x86)
  * to host BPF programs.
  */
-#ifdef CONFIG_TRANSPARENT_HUGEPAGE
-#define BPF_PROG_PACK_SIZE     HPAGE_PMD_SIZE
-#else
-#define BPF_PROG_PACK_SIZE     PAGE_SIZE
-#endif
 #define BPF_PROG_CHUNK_SHIFT   6
 #define BPF_PROG_CHUNK_SIZE    (1 << BPF_PROG_CHUNK_SHIFT)
 #define BPF_PROG_CHUNK_MASK    (~(BPF_PROG_CHUNK_SIZE - 1))
-#define BPF_PROG_CHUNK_COUNT   (BPF_PROG_PACK_SIZE / BPF_PROG_CHUNK_SIZE)
 
 struct bpf_prog_pack {
        struct list_head list;
        unsigned long bitmap[];
 };
 
-#define BPF_PROG_MAX_PACK_PROG_SIZE    BPF_PROG_PACK_SIZE
 #define BPF_PROG_SIZE_TO_NBITS(size)   (round_up(size, BPF_PROG_CHUNK_SIZE) / BPF_PROG_CHUNK_SIZE)
 
+static size_t bpf_prog_pack_size = -1;
+
+static int bpf_prog_chunk_count(void)
+{
+       WARN_ON_ONCE(bpf_prog_pack_size == -1);
+       return bpf_prog_pack_size / BPF_PROG_CHUNK_SIZE;
+}
+
 static DEFINE_MUTEX(pack_mutex);
 static LIST_HEAD(pack_list);
 
+static size_t select_bpf_prog_pack_size(void)
+{
+       size_t size;
+       void *ptr;
+
+       size = PMD_SIZE * num_online_nodes();
+       ptr = module_alloc(size);
+
+       /* Test whether we can get huge pages. If not just use PAGE_SIZE
+        * packs.
+        */
+       if (!ptr || !is_vm_area_hugepages(ptr))
+               size = PAGE_SIZE;
+
+       vfree(ptr);
+       return size;
+}
+
 static struct bpf_prog_pack *alloc_new_pack(void)
 {
        struct bpf_prog_pack *pack;
 
-       pack = kzalloc(sizeof(*pack) + BITS_TO_BYTES(BPF_PROG_CHUNK_COUNT), GFP_KERNEL);
+       pack = kzalloc(struct_size(pack, bitmap, BITS_TO_LONGS(bpf_prog_chunk_count())),
+                      GFP_KERNEL);
        if (!pack)
                return NULL;
-       pack->ptr = module_alloc(BPF_PROG_PACK_SIZE);
+       pack->ptr = module_alloc(bpf_prog_pack_size);
        if (!pack->ptr) {
                kfree(pack);
                return NULL;
        }
-       bitmap_zero(pack->bitmap, BPF_PROG_PACK_SIZE / BPF_PROG_CHUNK_SIZE);
+       bitmap_zero(pack->bitmap, bpf_prog_pack_size / BPF_PROG_CHUNK_SIZE);
        list_add_tail(&pack->list, &pack_list);
 
        set_vm_flush_reset_perms(pack->ptr);
-       set_memory_ro((unsigned long)pack->ptr, BPF_PROG_PACK_SIZE / PAGE_SIZE);
-       set_memory_x((unsigned long)pack->ptr, BPF_PROG_PACK_SIZE / PAGE_SIZE);
+       set_memory_ro((unsigned long)pack->ptr, bpf_prog_pack_size / PAGE_SIZE);
+       set_memory_x((unsigned long)pack->ptr, bpf_prog_pack_size / PAGE_SIZE);
        return pack;
 }
 
        unsigned long pos;
        void *ptr = NULL;
 
-       if (size > BPF_PROG_MAX_PACK_PROG_SIZE) {
+       mutex_lock(&pack_mutex);
+       if (bpf_prog_pack_size == -1)
+               bpf_prog_pack_size = select_bpf_prog_pack_size();
+
+       if (size > bpf_prog_pack_size) {
                size = round_up(size, PAGE_SIZE);
                ptr = module_alloc(size);
                if (ptr) {
                        set_memory_ro((unsigned long)ptr, size / PAGE_SIZE);
                        set_memory_x((unsigned long)ptr, size / PAGE_SIZE);
                }
-               return ptr;
+               goto out;
        }
-       mutex_lock(&pack_mutex);
        list_for_each_entry(pack, &pack_list, list) {
-               pos = bitmap_find_next_zero_area(pack->bitmap, BPF_PROG_CHUNK_COUNT, 0,
+               pos = bitmap_find_next_zero_area(pack->bitmap, bpf_prog_chunk_count(), 0,
                                                 nbits, 0);
-               if (pos < BPF_PROG_CHUNK_COUNT)
+               if (pos < bpf_prog_chunk_count())
                        goto found_free_area;
        }
 
        unsigned long pos;
        void *pack_ptr;
 
-       if (hdr->size > BPF_PROG_MAX_PACK_PROG_SIZE) {
+       mutex_lock(&pack_mutex);
+       if (hdr->size > bpf_prog_pack_size) {
                module_memfree(hdr);
-               return;
+               goto out;
        }
 
-       pack_ptr = (void *)((unsigned long)hdr & ~(BPF_PROG_PACK_SIZE - 1));
-       mutex_lock(&pack_mutex);
+       pack_ptr = (void *)((unsigned long)hdr & ~(bpf_prog_pack_size - 1));
 
        list_for_each_entry(tmp, &pack_list, list) {
                if (tmp->ptr == pack_ptr) {
        pos = ((unsigned long)hdr - (unsigned long)pack_ptr) >> BPF_PROG_CHUNK_SHIFT;
 
        bitmap_clear(pack->bitmap, pos, nbits);
-       if (bitmap_find_next_zero_area(pack->bitmap, BPF_PROG_CHUNK_COUNT, 0,
-                                      BPF_PROG_CHUNK_COUNT, 0) == 0) {
+       if (bitmap_find_next_zero_area(pack->bitmap, bpf_prog_chunk_count(), 0,
+                                      bpf_prog_chunk_count(), 0) == 0) {
                list_del(&pack->list);
                module_memfree(pack->ptr);
                kfree(pack);