struct bpf_local_storage_elem *
 bpf_selem_alloc(struct bpf_local_storage_map *smap, void *owner, void *value,
-               bool charge_mem);
+               bool charge_mem, gfp_t gfp_flags);
 
 int
 bpf_local_storage_alloc(void *owner,
                        struct bpf_local_storage_map *smap,
-                       struct bpf_local_storage_elem *first_selem);
+                       struct bpf_local_storage_elem *first_selem,
+                       gfp_t gfp_flags);
 
 struct bpf_local_storage_data *
 bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
-                        void *value, u64 map_flags);
+                        void *value, u64 map_flags, gfp_t gfp_flags);
 
 void bpf_local_storage_free_rcu(struct rcu_head *rcu);
 
 
 
        sdata = bpf_local_storage_update(f->f_inode,
                                         (struct bpf_local_storage_map *)map,
-                                        value, map_flags);
+                                        value, map_flags, GFP_ATOMIC);
        fput(f);
        return PTR_ERR_OR_ZERO(sdata);
 }
        return err;
 }
 
-BPF_CALL_4(bpf_inode_storage_get, struct bpf_map *, map, struct inode *, inode,
-          void *, value, u64, flags)
+/* *gfp_flags* is a hidden argument provided by the verifier */
+BPF_CALL_5(bpf_inode_storage_get, struct bpf_map *, map, struct inode *, inode,
+          void *, value, u64, flags, gfp_t, gfp_flags)
 {
        struct bpf_local_storage_data *sdata;
 
        if (flags & BPF_LOCAL_STORAGE_GET_F_CREATE) {
                sdata = bpf_local_storage_update(
                        inode, (struct bpf_local_storage_map *)map, value,
-                       BPF_NOEXIST);
+                       BPF_NOEXIST, gfp_flags);
                return IS_ERR(sdata) ? (unsigned long)NULL :
                                             (unsigned long)sdata->data;
        }
 
 
 struct bpf_local_storage_elem *
 bpf_selem_alloc(struct bpf_local_storage_map *smap, void *owner,
-               void *value, bool charge_mem)
+               void *value, bool charge_mem, gfp_t gfp_flags)
 {
        struct bpf_local_storage_elem *selem;
 
                return NULL;
 
        selem = bpf_map_kzalloc(&smap->map, smap->elem_size,
-                               GFP_ATOMIC | __GFP_NOWARN);
+                               gfp_flags | __GFP_NOWARN);
        if (selem) {
                if (value)
                        memcpy(SDATA(selem)->data, value, smap->map.value_size);
 
 int bpf_local_storage_alloc(void *owner,
                            struct bpf_local_storage_map *smap,
-                           struct bpf_local_storage_elem *first_selem)
+                           struct bpf_local_storage_elem *first_selem,
+                           gfp_t gfp_flags)
 {
        struct bpf_local_storage *prev_storage, *storage;
        struct bpf_local_storage **owner_storage_ptr;
                return err;
 
        storage = bpf_map_kzalloc(&smap->map, sizeof(*storage),
-                                 GFP_ATOMIC | __GFP_NOWARN);
+                                 gfp_flags | __GFP_NOWARN);
        if (!storage) {
                err = -ENOMEM;
                goto uncharge;
  */
 struct bpf_local_storage_data *
 bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
-                        void *value, u64 map_flags)
+                        void *value, u64 map_flags, gfp_t gfp_flags)
 {
        struct bpf_local_storage_data *old_sdata = NULL;
-       struct bpf_local_storage_elem *selem;
+       struct bpf_local_storage_elem *selem = NULL;
        struct bpf_local_storage *local_storage;
        unsigned long flags;
        int err;
                     !map_value_has_spin_lock(&smap->map)))
                return ERR_PTR(-EINVAL);
 
+       if (gfp_flags == GFP_KERNEL && (map_flags & ~BPF_F_LOCK) != BPF_NOEXIST)
+               return ERR_PTR(-EINVAL);
+
        local_storage = rcu_dereference_check(*owner_storage(smap, owner),
                                              bpf_rcu_lock_held());
        if (!local_storage || hlist_empty(&local_storage->list)) {
                if (err)
                        return ERR_PTR(err);
 
-               selem = bpf_selem_alloc(smap, owner, value, true);
+               selem = bpf_selem_alloc(smap, owner, value, true, gfp_flags);
                if (!selem)
                        return ERR_PTR(-ENOMEM);
 
-               err = bpf_local_storage_alloc(owner, smap, selem);
+               err = bpf_local_storage_alloc(owner, smap, selem, gfp_flags);
                if (err) {
                        kfree(selem);
                        mem_uncharge(smap, owner, smap->elem_size);
                }
        }
 
+       if (gfp_flags == GFP_KERNEL) {
+               selem = bpf_selem_alloc(smap, owner, value, true, gfp_flags);
+               if (!selem)
+                       return ERR_PTR(-ENOMEM);
+       }
+
        raw_spin_lock_irqsave(&local_storage->lock, flags);
 
        /* Recheck local_storage->list under local_storage->lock */
                goto unlock;
        }
 
-       /* local_storage->lock is held.  Hence, we are sure
-        * we can unlink and uncharge the old_sdata successfully
-        * later.  Hence, instead of charging the new selem now
-        * and then uncharge the old selem later (which may cause
-        * a potential but unnecessary charge failure),  avoid taking
-        * a charge at all here (the "!old_sdata" check) and the
-        * old_sdata will not be uncharged later during
-        * bpf_selem_unlink_storage_nolock().
-        */
-       selem = bpf_selem_alloc(smap, owner, value, !old_sdata);
-       if (!selem) {
-               err = -ENOMEM;
-               goto unlock_err;
+       if (gfp_flags != GFP_KERNEL) {
+               /* local_storage->lock is held.  Hence, we are sure
+                * we can unlink and uncharge the old_sdata successfully
+                * later.  Hence, instead of charging the new selem now
+                * and then uncharge the old selem later (which may cause
+                * a potential but unnecessary charge failure),  avoid taking
+                * a charge at all here (the "!old_sdata" check) and the
+                * old_sdata will not be uncharged later during
+                * bpf_selem_unlink_storage_nolock().
+                */
+               selem = bpf_selem_alloc(smap, owner, value, !old_sdata, gfp_flags);
+               if (!selem) {
+                       err = -ENOMEM;
+                       goto unlock_err;
+               }
        }
 
        /* First, link the new selem to the map */
 
 unlock_err:
        raw_spin_unlock_irqrestore(&local_storage->lock, flags);
+       if (selem) {
+               mem_uncharge(smap, owner, smap->elem_size);
+               kfree(selem);
+       }
        return ERR_PTR(err);
 }
 
 
 
        bpf_task_storage_lock();
        sdata = bpf_local_storage_update(
-               task, (struct bpf_local_storage_map *)map, value, map_flags);
+               task, (struct bpf_local_storage_map *)map, value, map_flags,
+               GFP_ATOMIC);
        bpf_task_storage_unlock();
 
        err = PTR_ERR_OR_ZERO(sdata);
        return err;
 }
 
-BPF_CALL_4(bpf_task_storage_get, struct bpf_map *, map, struct task_struct *,
-          task, void *, value, u64, flags)
+/* *gfp_flags* is a hidden argument provided by the verifier */
+BPF_CALL_5(bpf_task_storage_get, struct bpf_map *, map, struct task_struct *,
+          task, void *, value, u64, flags, gfp_t, gfp_flags)
 {
        struct bpf_local_storage_data *sdata;
 
            (flags & BPF_LOCAL_STORAGE_GET_F_CREATE))
                sdata = bpf_local_storage_update(
                        task, (struct bpf_local_storage_map *)map, value,
-                       BPF_NOEXIST);
+                       BPF_NOEXIST, gfp_flags);
 
 unlock:
        bpf_task_storage_unlock();
 
                        goto patch_call_imm;
                }
 
+               if (insn->imm == BPF_FUNC_task_storage_get ||
+                   insn->imm == BPF_FUNC_sk_storage_get ||
+                   insn->imm == BPF_FUNC_inode_storage_get) {
+                       if (env->prog->aux->sleepable)
+                               insn_buf[0] = BPF_MOV64_IMM(BPF_REG_5, (__s32)GFP_KERNEL);
+                       else
+                               insn_buf[0] = BPF_MOV64_IMM(BPF_REG_5, (__s32)GFP_ATOMIC);
+                       insn_buf[1] = *insn;
+                       cnt = 2;
+
+                       new_prog = bpf_patch_insn_data(env, i + delta, insn_buf, cnt);
+                       if (!new_prog)
+                               return -ENOMEM;
+
+                       delta += cnt - 1;
+                       env->prog = prog = new_prog;
+                       insn = new_prog->insnsi + i + delta;
+                       goto patch_call_imm;
+               }
+
                /* BPF_EMIT_CALL() assumptions in some of the map_gen_lookup
                 * and other inlining handlers are currently limited to 64 bit
                 * only.
 
        if (sock) {
                sdata = bpf_local_storage_update(
                        sock->sk, (struct bpf_local_storage_map *)map, value,
-                       map_flags);
+                       map_flags, GFP_ATOMIC);
                sockfd_put(sock);
                return PTR_ERR_OR_ZERO(sdata);
        }
 {
        struct bpf_local_storage_elem *copy_selem;
 
-       copy_selem = bpf_selem_alloc(smap, newsk, NULL, true);
+       copy_selem = bpf_selem_alloc(smap, newsk, NULL, true, GFP_ATOMIC);
        if (!copy_selem)
                return NULL;
 
                        bpf_selem_link_map(smap, copy_selem);
                        bpf_selem_link_storage_nolock(new_sk_storage, copy_selem);
                } else {
-                       ret = bpf_local_storage_alloc(newsk, smap, copy_selem);
+                       ret = bpf_local_storage_alloc(newsk, smap, copy_selem, GFP_ATOMIC);
                        if (ret) {
                                kfree(copy_selem);
                                atomic_sub(smap->elem_size,
        return ret;
 }
 
-BPF_CALL_4(bpf_sk_storage_get, struct bpf_map *, map, struct sock *, sk,
-          void *, value, u64, flags)
+/* *gfp_flags* is a hidden argument provided by the verifier */
+BPF_CALL_5(bpf_sk_storage_get, struct bpf_map *, map, struct sock *, sk,
+          void *, value, u64, flags, gfp_t, gfp_flags)
 {
        struct bpf_local_storage_data *sdata;
 
            refcount_inc_not_zero(&sk->sk_refcnt)) {
                sdata = bpf_local_storage_update(
                        sk, (struct bpf_local_storage_map *)map, value,
-                       BPF_NOEXIST);
+                       BPF_NOEXIST, gfp_flags);
                /* sk must be a fullsock (guaranteed by verifier),
                 * so sock_gen_put() is unnecessary.
                 */
        return false;
 }
 
-BPF_CALL_4(bpf_sk_storage_get_tracing, struct bpf_map *, map, struct sock *, sk,
-          void *, value, u64, flags)
+/* *gfp_flags* is a hidden argument provided by the verifier */
+BPF_CALL_5(bpf_sk_storage_get_tracing, struct bpf_map *, map, struct sock *, sk,
+          void *, value, u64, flags, gfp_t, gfp_flags)
 {
        WARN_ON_ONCE(!bpf_rcu_lock_held());
        if (in_hardirq() || in_nmi())
                return (unsigned long)NULL;
 
-       return (unsigned long)____bpf_sk_storage_get(map, sk, value, flags);
+       return (unsigned long)____bpf_sk_storage_get(map, sk, value, flags,
+                                                    gfp_flags);
 }
 
 BPF_CALL_2(bpf_sk_storage_delete_tracing, struct bpf_map *, map,