static DEFINE_PER_CPU(struct page_frag_cache, netdev_alloc_cache);
 static DEFINE_PER_CPU(struct napi_alloc_cache, napi_alloc_cache);
 
-static void *__netdev_alloc_frag(unsigned int fragsz, gfp_t gfp_mask)
+static void *__napi_alloc_frag(unsigned int fragsz, gfp_t gfp_mask)
 {
-       struct page_frag_cache *nc;
-       unsigned long flags;
-       void *data;
+       struct napi_alloc_cache *nc = this_cpu_ptr(&napi_alloc_cache);
 
-       local_irq_save(flags);
-       nc = this_cpu_ptr(&netdev_alloc_cache);
-       data = page_frag_alloc(nc, fragsz, gfp_mask);
-       local_irq_restore(flags);
-       return data;
+       return page_frag_alloc(&nc->page, fragsz, gfp_mask);
 }
 
+void *napi_alloc_frag(unsigned int fragsz)
+{
+       fragsz = SKB_DATA_ALIGN(fragsz);
+
+       return __napi_alloc_frag(fragsz, GFP_ATOMIC);
+}
+EXPORT_SYMBOL(napi_alloc_frag);
+
 /**
  * netdev_alloc_frag - allocate a page fragment
  * @fragsz: fragment size
  */
 void *netdev_alloc_frag(unsigned int fragsz)
 {
-       fragsz = SKB_DATA_ALIGN(fragsz);
-
-       return __netdev_alloc_frag(fragsz, GFP_ATOMIC);
-}
-EXPORT_SYMBOL(netdev_alloc_frag);
-
-static void *__napi_alloc_frag(unsigned int fragsz, gfp_t gfp_mask)
-{
-       struct napi_alloc_cache *nc = this_cpu_ptr(&napi_alloc_cache);
-
-       return page_frag_alloc(&nc->page, fragsz, gfp_mask);
-}
+       struct page_frag_cache *nc;
+       void *data;
 
-void *napi_alloc_frag(unsigned int fragsz)
-{
        fragsz = SKB_DATA_ALIGN(fragsz);
-
-       return __napi_alloc_frag(fragsz, GFP_ATOMIC);
+       if (in_irq() || irqs_disabled()) {
+               nc = this_cpu_ptr(&netdev_alloc_cache);
+               data = page_frag_alloc(nc, fragsz, GFP_ATOMIC);
+       } else {
+               local_bh_disable();
+               data = __napi_alloc_frag(fragsz, GFP_ATOMIC);
+               local_bh_enable();
+       }
+       return data;
 }
-EXPORT_SYMBOL(napi_alloc_frag);
+EXPORT_SYMBOL(netdev_alloc_frag);
 
 /**
  *     __netdev_alloc_skb - allocate an skbuff for rx on a specific device