diff mbox series

[bpf-next,v3,3/5] bpf: Mark each subprog with proper pstack states

Message ID 20240926234521.1770481-1-yonghong.song@linux.dev (mailing list archive)
State Superseded
Delegated to: BPF
Headers show
Series bpf: Support private stack for bpf progs | expand

Checks

Context Check Description
bpf/vmtest-bpf-next-PR success PR summary
netdev/series_format success Posting correctly formatted
netdev/tree_selection success Clearly marked for bpf-next, async
netdev/ynl success Generated files up to date; no warnings/errors; no diff in generated;
netdev/fixes_present success Fixes tag not required for -next series
netdev/header_inline success No static functions without inline keyword in header files
netdev/build_32bit success Errors and warnings before: 213 this patch: 213
netdev/build_tools success Errors and warnings before: 0 this patch: 0
netdev/cc_maintainers warning 8 maintainers not CCed: sdf@fomichev.me eddyz87@gmail.com haoluo@google.com jolsa@kernel.org song@kernel.org kpsingh@kernel.org martin.lau@linux.dev john.fastabend@gmail.com
netdev/build_clang success Errors and warnings before: 272 this patch: 272
netdev/verify_signedoff success Signed-off-by tag matches author and committer
netdev/deprecated_api success None detected
netdev/check_selftest success No net selftest shell script
netdev/verify_fixes success No Fixes tag
netdev/build_allmodconfig_warn success Errors and warnings before: 6964 this patch: 6964
netdev/checkpatch warning WARNING: line length of 116 exceeds 80 columns WARNING: line length of 81 exceeds 80 columns WARNING: line length of 91 exceeds 80 columns WARNING: line length of 95 exceeds 80 columns WARNING: line length of 96 exceeds 80 columns
netdev/build_clang_rust success No Rust files in patch. Skipping build
netdev/kdoc success Errors and warnings before: 6 this patch: 6
netdev/source_inline success Was 0 now: 0
bpf/vmtest-bpf-next-VM_Test-0 success Logs for Lint
bpf/vmtest-bpf-next-VM_Test-1 success Logs for ShellCheck
bpf/vmtest-bpf-next-VM_Test-2 success Logs for Unittests
bpf/vmtest-bpf-next-VM_Test-3 success Logs for Validate matrix.py
bpf/vmtest-bpf-next-VM_Test-5 success Logs for aarch64-gcc / build-release
bpf/vmtest-bpf-next-VM_Test-4 success Logs for aarch64-gcc / build / build for aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-10 success Logs for aarch64-gcc / veristat
bpf/vmtest-bpf-next-VM_Test-12 success Logs for s390x-gcc / build-release
bpf/vmtest-bpf-next-VM_Test-6 success Logs for aarch64-gcc / test (test_maps, false, 360) / test_maps on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-9 success Logs for aarch64-gcc / test (test_verifier, false, 360) / test_verifier on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-7 success Logs for aarch64-gcc / test (test_progs, false, 360) / test_progs on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-8 success Logs for aarch64-gcc / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-11 success Logs for s390x-gcc / build / build for s390x with gcc
bpf/vmtest-bpf-next-VM_Test-19 success Logs for x86_64-gcc / build-release
bpf/vmtest-bpf-next-VM_Test-27 success Logs for x86_64-llvm-17 / build / build for x86_64 with llvm-17
bpf/vmtest-bpf-next-VM_Test-28 success Logs for x86_64-llvm-17 / build-release / build for x86_64 with llvm-17-O2
bpf/vmtest-bpf-next-VM_Test-33 success Logs for x86_64-llvm-17 / veristat
bpf/vmtest-bpf-next-VM_Test-17 success Logs for set-matrix
bpf/vmtest-bpf-next-VM_Test-16 success Logs for s390x-gcc / veristat
bpf/vmtest-bpf-next-VM_Test-35 success Logs for x86_64-llvm-18 / build-release / build for x86_64 with llvm-18-O2
bpf/vmtest-bpf-next-VM_Test-18 success Logs for x86_64-gcc / build / build for x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-15 success Logs for s390x-gcc / test (test_verifier, false, 360) / test_verifier on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-34 success Logs for x86_64-llvm-18 / build / build for x86_64 with llvm-18
bpf/vmtest-bpf-next-VM_Test-41 success Logs for x86_64-llvm-18 / veristat
bpf/vmtest-bpf-next-VM_Test-13 success Logs for s390x-gcc / test (test_progs, false, 360) / test_progs on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-14 success Logs for s390x-gcc / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-25 success Logs for x86_64-gcc / test (test_verifier, false, 360) / test_verifier on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-20 success Logs for x86_64-gcc / test (test_maps, false, 360) / test_maps on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-21 success Logs for x86_64-gcc / test (test_progs, false, 360) / test_progs on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-22 success Logs for x86_64-gcc / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-23 success Logs for x86_64-gcc / test (test_progs_no_alu32_parallel, true, 30) / test_progs_no_alu32_parallel on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-24 success Logs for x86_64-gcc / test (test_progs_parallel, true, 30) / test_progs_parallel on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-26 success Logs for x86_64-gcc / veristat / veristat on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-29 success Logs for x86_64-llvm-17 / test (test_maps, false, 360) / test_maps on x86_64 with llvm-17
bpf/vmtest-bpf-next-VM_Test-30 success Logs for x86_64-llvm-17 / test (test_progs, false, 360) / test_progs on x86_64 with llvm-17
bpf/vmtest-bpf-next-VM_Test-31 success Logs for x86_64-llvm-17 / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on x86_64 with llvm-17
bpf/vmtest-bpf-next-VM_Test-32 success Logs for x86_64-llvm-17 / test (test_verifier, false, 360) / test_verifier on x86_64 with llvm-17
bpf/vmtest-bpf-next-VM_Test-36 success Logs for x86_64-llvm-18 / test (test_maps, false, 360) / test_maps on x86_64 with llvm-18
bpf/vmtest-bpf-next-VM_Test-37 success Logs for x86_64-llvm-18 / test (test_progs, false, 360) / test_progs on x86_64 with llvm-18
bpf/vmtest-bpf-next-VM_Test-38 success Logs for x86_64-llvm-18 / test (test_progs_cpuv4, false, 360) / test_progs_cpuv4 on x86_64 with llvm-18
bpf/vmtest-bpf-next-VM_Test-39 success Logs for x86_64-llvm-18 / test (test_progs_no_alu32, false, 360) / test_progs_no_alu32 on x86_64 with llvm-18
bpf/vmtest-bpf-next-VM_Test-40 success Logs for x86_64-llvm-18 / test (test_verifier, false, 360) / test_verifier on x86_64 with llvm-18

Commit Message

Yonghong Song Sept. 26, 2024, 11:45 p.m. UTC
Three private stack states are used to direct jit action:
  PSTACK_TREE_NO:       do not use private stack
  PSTACK_TREE_INTERNAL: adjust frame pointer address (similar to normal stack)
  PSTACK_TREE_ROOT:     set the frame pointer

Note that for subtree root, even if the root bpf_prog stack size is 0,
PSTACK_TREE_INTERNAL is still used. This is for bpf exception handling.
More details can be found in subsequent jit support and selftest patches.

Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
---
 include/linux/bpf.h   |  9 +++++++++
 kernel/bpf/core.c     | 19 +++++++++++++++++++
 kernel/bpf/verifier.c | 30 ++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+)

Comments

Alexei Starovoitov Sept. 30, 2024, 2:49 p.m. UTC | #1
On Thu, Sep 26, 2024 at 4:45 PM Yonghong Song <yonghong.song@linux.dev> wrote:
>
> Three private stack states are used to direct jit action:
>   PSTACK_TREE_NO:       do not use private stack
>   PSTACK_TREE_INTERNAL: adjust frame pointer address (similar to normal stack)
>   PSTACK_TREE_ROOT:     set the frame pointer
>
> Note that for subtree root, even if the root bpf_prog stack size is 0,
> PSTACK_TREE_INTERNAL is still used. This is for bpf exception handling.
> More details can be found in subsequent jit support and selftest patches.
>
> Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
> ---
>  include/linux/bpf.h   |  9 +++++++++
>  kernel/bpf/core.c     | 19 +++++++++++++++++++
>  kernel/bpf/verifier.c | 30 ++++++++++++++++++++++++++++++
>  3 files changed, 58 insertions(+)
>
> diff --git a/include/linux/bpf.h b/include/linux/bpf.h
> index 156b9516d9f6..8f02d11bd408 100644
> --- a/include/linux/bpf.h
> +++ b/include/linux/bpf.h
> @@ -1550,6 +1550,12 @@ struct bpf_prog_aux {
>         };
>  };
>
> +enum bpf_pstack_state {
> +       PSTACK_TREE_NO,
> +       PSTACK_TREE_INTERNAL,
> +       PSTACK_TREE_ROOT,
> +};

The names could be improved and 'state' doesn't quite fit imo.
How about:
enum bpf_priv_stack_mode {
   NO_PRIV_STACK,
   PRIV_STACK_SUB_PROG,
   PRIV_STACK_MAIN_PROG,
};

> +
>  struct bpf_prog {
>         u16                     pages;          /* Number of allocated pages */
>         u16                     jited:1,        /* Is our filter JIT'ed? */
> @@ -1570,15 +1576,18 @@ struct bpf_prog {
>                                 pstack_eligible:1; /* Candidate for private stacks */
>         enum bpf_prog_type      type;           /* Type of BPF program */
>         enum bpf_attach_type    expected_attach_type; /* For some prog types */
> +       enum bpf_pstack_state   pstack:2;       /* Private stack state */
>         u32                     len;            /* Number of filter blocks */
>         u32                     jited_len;      /* Size of jited insns in bytes */
>         u8                      tag[BPF_TAG_SIZE];
> +       u16                     subtree_stack_depth; /* Subtree stack depth if PSTACK_TREE_ROOT prog, 0 otherwise */

All the extra vars can be in prog->aux.
No need to put them in struct bpf_prog.

>         struct bpf_prog_stats __percpu *stats;
>         int __percpu            *active;
>         unsigned int            (*bpf_func)(const void *ctx,
>                                             const struct bpf_insn *insn);
>         struct bpf_prog_aux     *aux;           /* Auxiliary fields */
>         struct sock_fprog_kern  *orig_prog;     /* Original BPF program */
> +       void __percpu           *private_stack_ptr;

same as this one. prog->aux should be fine.

>         /* Instructions for interpreter */
>         union {
>                 DECLARE_FLEX_ARRAY(struct sock_filter, insns);
> diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
> index 0727fff6de0e..d6eb052f6631 100644
> --- a/kernel/bpf/core.c
> +++ b/kernel/bpf/core.c
> @@ -1239,6 +1239,7 @@ void __weak bpf_jit_free(struct bpf_prog *fp)
>                 struct bpf_binary_header *hdr = bpf_jit_binary_hdr(fp);
>
>                 bpf_jit_binary_free(hdr);
> +               free_percpu(fp->private_stack_ptr);
>                 WARN_ON_ONCE(!bpf_prog_kallsyms_verify_off(fp));
>         }
>
> @@ -2420,6 +2421,24 @@ struct bpf_prog *bpf_prog_select_runtime(struct bpf_prog *fp, int *err)
>                 if (*err)
>                         return fp;
>
> +               if (fp->pstack_eligible) {
> +                       if (!fp->aux->stack_depth) {
> +                               fp->pstack = PSTACK_TREE_NO;
> +                       } else {
> +                               void __percpu *private_stack_ptr;
> +
> +                               fp->pstack = PSTACK_TREE_ROOT;
> +                               private_stack_ptr =
> +                                       __alloc_percpu_gfp(fp->aux->stack_depth, 8, GFP_KERNEL);
> +                               if (!private_stack_ptr) {
> +                                       *err = -ENOMEM;
> +                                       return fp;
> +                               }
> +                               fp->subtree_stack_depth = fp->aux->stack_depth;
> +                               fp->private_stack_ptr = private_stack_ptr;
> +                       }
> +               }
> +
>                 fp = bpf_int_jit_compile(fp);
>                 bpf_prog_jit_attempt_done(fp);
>                 if (!fp->jited && jit_needed) {
> diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
> index 69e17cb22037..9d093e2013ca 100644
> --- a/kernel/bpf/verifier.c
> +++ b/kernel/bpf/verifier.c
> @@ -20060,6 +20060,7 @@ static int jit_subprogs(struct bpf_verifier_env *env)
>  {
>         struct bpf_prog *prog = env->prog, **func, *tmp;
>         int i, j, subprog_start, subprog_end = 0, len, subprog;
> +       int subtree_top_idx, subtree_stack_depth;
>         struct bpf_map *map_ptr;
>         struct bpf_insn *insn;
>         void *old_bpf_func;
> @@ -20138,6 +20139,35 @@ static int jit_subprogs(struct bpf_verifier_env *env)
>                 func[i]->is_func = 1;
>                 func[i]->sleepable = prog->sleepable;
>                 func[i]->aux->func_idx = i;
> +
> +               subtree_top_idx = env->subprog_info[i].subtree_top_idx;
> +               if (env->subprog_info[subtree_top_idx].pstack_eligible) {
> +                       if (subtree_top_idx == i)
> +                               func[i]->subtree_stack_depth =
> +                                       env->subprog_info[i].subtree_stack_depth;
> +
> +                       subtree_stack_depth = func[i]->subtree_stack_depth;
> +                       if (subtree_top_idx != i) {
> +                               if (env->subprog_info[subtree_top_idx].subtree_stack_depth)
> +                                       func[i]->pstack = PSTACK_TREE_INTERNAL;
> +                               else
> +                                       func[i]->pstack = PSTACK_TREE_NO;
> +                       } else if (!subtree_stack_depth) {
> +                               func[i]->pstack = PSTACK_TREE_INTERNAL;
> +                       } else {
> +                               void __percpu *private_stack_ptr;
> +
> +                               func[i]->pstack = PSTACK_TREE_ROOT;
> +                               private_stack_ptr =
> +                                       __alloc_percpu_gfp(subtree_stack_depth, 8, GFP_KERNEL);
> +                               if (!private_stack_ptr) {
> +                                       err = -ENOMEM;
> +                                       goto out_free;
> +                               }
> +                               func[i]->private_stack_ptr = private_stack_ptr;
> +                       }
> +               }
> +
>                 /* Below members will be freed only at prog->aux */
>                 func[i]->aux->btf = prog->aux->btf;
>                 func[i]->aux->func_info = prog->aux->func_info;
> --
> 2.43.5
>
Yonghong Song Sept. 30, 2024, 4:26 p.m. UTC | #2
On 9/30/24 7:49 AM, Alexei Starovoitov wrote:
> On Thu, Sep 26, 2024 at 4:45 PM Yonghong Song <yonghong.song@linux.dev> wrote:
>> Three private stack states are used to direct jit action:
>>    PSTACK_TREE_NO:       do not use private stack
>>    PSTACK_TREE_INTERNAL: adjust frame pointer address (similar to normal stack)
>>    PSTACK_TREE_ROOT:     set the frame pointer
>>
>> Note that for subtree root, even if the root bpf_prog stack size is 0,
>> PSTACK_TREE_INTERNAL is still used. This is for bpf exception handling.
>> More details can be found in subsequent jit support and selftest patches.
>>
>> Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
>> ---
>>   include/linux/bpf.h   |  9 +++++++++
>>   kernel/bpf/core.c     | 19 +++++++++++++++++++
>>   kernel/bpf/verifier.c | 30 ++++++++++++++++++++++++++++++
>>   3 files changed, 58 insertions(+)
>>
>> diff --git a/include/linux/bpf.h b/include/linux/bpf.h
>> index 156b9516d9f6..8f02d11bd408 100644
>> --- a/include/linux/bpf.h
>> +++ b/include/linux/bpf.h
>> @@ -1550,6 +1550,12 @@ struct bpf_prog_aux {
>>          };
>>   };
>>
>> +enum bpf_pstack_state {
>> +       PSTACK_TREE_NO,
>> +       PSTACK_TREE_INTERNAL,
>> +       PSTACK_TREE_ROOT,
>> +};
> The names could be improved and 'state' doesn't quite fit imo.
> How about:
> enum bpf_priv_stack_mode {
>     NO_PRIV_STACK,
>     PRIV_STACK_SUB_PROG,
>     PRIV_STACK_MAIN_PROG,
> };

Since we agreed to use priv_stack instead of pstack. The above
names make sense. Will change.

>
>> +
>>   struct bpf_prog {
>>          u16                     pages;          /* Number of allocated pages */
>>          u16                     jited:1,        /* Is our filter JIT'ed? */
>> @@ -1570,15 +1576,18 @@ struct bpf_prog {
>>                                  pstack_eligible:1; /* Candidate for private stacks */
>>          enum bpf_prog_type      type;           /* Type of BPF program */
>>          enum bpf_attach_type    expected_attach_type; /* For some prog types */
>> +       enum bpf_pstack_state   pstack:2;       /* Private stack state */
>>          u32                     len;            /* Number of filter blocks */
>>          u32                     jited_len;      /* Size of jited insns in bytes */
>>          u8                      tag[BPF_TAG_SIZE];
>> +       u16                     subtree_stack_depth; /* Subtree stack depth if PSTACK_TREE_ROOT prog, 0 otherwise */
> All the extra vars can be in prog->aux.
> No need to put them in struct bpf_prog.

Will do.

>
>>          struct bpf_prog_stats __percpu *stats;
>>          int __percpu            *active;
>>          unsigned int            (*bpf_func)(const void *ctx,
>>                                              const struct bpf_insn *insn);
>>          struct bpf_prog_aux     *aux;           /* Auxiliary fields */
>>          struct sock_fprog_kern  *orig_prog;     /* Original BPF program */
>> +       void __percpu           *private_stack_ptr;
> same as this one. prog->aux should be fine.

Will do.

>
>>          /* Instructions for interpreter */
>>          union {
>>                  DECLARE_FLEX_ARRAY(struct sock_filter, insns);
>> diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
>> index 0727fff6de0e..d6eb052f6631 100644
>> --- a/kernel/bpf/core.c
>> +++ b/kernel/bpf/core.c
>> @@ -1239,6 +1239,7 @@ void __weak bpf_jit_free(struct bpf_prog *fp)
>>                  struct bpf_binary_header *hdr = bpf_jit_binary_hdr(fp);
>>
>>                  bpf_jit_binary_free(hdr);
>> +               free_percpu(fp->private_stack_ptr);
>>                  WARN_ON_ONCE(!bpf_prog_kallsyms_verify_off(fp));
>>          }
>>
>> @@ -2420,6 +2421,24 @@ struct bpf_prog *bpf_prog_select_runtime(struct bpf_prog *fp, int *err)
>>                  if (*err)
>>                          return fp;
>>
>> +               if (fp->pstack_eligible) {
>> +                       if (!fp->aux->stack_depth) {
>> +                               fp->pstack = PSTACK_TREE_NO;
>> +                       } else {
>> +                               void __percpu *private_stack_ptr;
>> +
>> +                               fp->pstack = PSTACK_TREE_ROOT;
>> +                               private_stack_ptr =
>> +                                       __alloc_percpu_gfp(fp->aux->stack_depth, 8, GFP_KERNEL);
>> +                               if (!private_stack_ptr) {
>> +                                       *err = -ENOMEM;
>> +                                       return fp;
>> +                               }
>> +                               fp->subtree_stack_depth = fp->aux->stack_depth;
>> +                               fp->private_stack_ptr = private_stack_ptr;
>> +                       }
>> +               }
>> +
>>                  fp = bpf_int_jit_compile(fp);
>>                  bpf_prog_jit_attempt_done(fp);
>>                  if (!fp->jited && jit_needed) {
>> diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
>> index 69e17cb22037..9d093e2013ca 100644
>> --- a/kernel/bpf/verifier.c
>> +++ b/kernel/bpf/verifier.c
>> @@ -20060,6 +20060,7 @@ static int jit_subprogs(struct bpf_verifier_env *env)
>>   {
>>          struct bpf_prog *prog = env->prog, **func, *tmp;
>>          int i, j, subprog_start, subprog_end = 0, len, subprog;
>> +       int subtree_top_idx, subtree_stack_depth;
>>          struct bpf_map *map_ptr;
>>          struct bpf_insn *insn;
>>          void *old_bpf_func;
>> @@ -20138,6 +20139,35 @@ static int jit_subprogs(struct bpf_verifier_env *env)
>>                  func[i]->is_func = 1;
>>                  func[i]->sleepable = prog->sleepable;
>>                  func[i]->aux->func_idx = i;
>> +
>> +               subtree_top_idx = env->subprog_info[i].subtree_top_idx;
>> +               if (env->subprog_info[subtree_top_idx].pstack_eligible) {
>> +                       if (subtree_top_idx == i)
>> +                               func[i]->subtree_stack_depth =
>> +                                       env->subprog_info[i].subtree_stack_depth;
>> +
>> +                       subtree_stack_depth = func[i]->subtree_stack_depth;
>> +                       if (subtree_top_idx != i) {
>> +                               if (env->subprog_info[subtree_top_idx].subtree_stack_depth)
>> +                                       func[i]->pstack = PSTACK_TREE_INTERNAL;
>> +                               else
>> +                                       func[i]->pstack = PSTACK_TREE_NO;
>> +                       } else if (!subtree_stack_depth) {
>> +                               func[i]->pstack = PSTACK_TREE_INTERNAL;
>> +                       } else {
>> +                               void __percpu *private_stack_ptr;
>> +
>> +                               func[i]->pstack = PSTACK_TREE_ROOT;
>> +                               private_stack_ptr =
>> +                                       __alloc_percpu_gfp(subtree_stack_depth, 8, GFP_KERNEL);
>> +                               if (!private_stack_ptr) {
>> +                                       err = -ENOMEM;
>> +                                       goto out_free;
>> +                               }
>> +                               func[i]->private_stack_ptr = private_stack_ptr;
>> +                       }
>> +               }
>> +
>>                  /* Below members will be freed only at prog->aux */
>>                  func[i]->aux->btf = prog->aux->btf;
>>                  func[i]->aux->func_info = prog->aux->func_info;
>> --
>> 2.43.5
>>
diff mbox series

Patch

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 156b9516d9f6..8f02d11bd408 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1550,6 +1550,12 @@  struct bpf_prog_aux {
 	};
 };
 
+enum bpf_pstack_state {
+	PSTACK_TREE_NO,
+	PSTACK_TREE_INTERNAL,
+	PSTACK_TREE_ROOT,
+};
+
 struct bpf_prog {
 	u16			pages;		/* Number of allocated pages */
 	u16			jited:1,	/* Is our filter JIT'ed? */
@@ -1570,15 +1576,18 @@  struct bpf_prog {
 				pstack_eligible:1; /* Candidate for private stacks */
 	enum bpf_prog_type	type;		/* Type of BPF program */
 	enum bpf_attach_type	expected_attach_type; /* For some prog types */
+	enum bpf_pstack_state	pstack:2;	/* Private stack state */
 	u32			len;		/* Number of filter blocks */
 	u32			jited_len;	/* Size of jited insns in bytes */
 	u8			tag[BPF_TAG_SIZE];
+	u16			subtree_stack_depth; /* Subtree stack depth if PSTACK_TREE_ROOT prog, 0 otherwise */
 	struct bpf_prog_stats __percpu *stats;
 	int __percpu		*active;
 	unsigned int		(*bpf_func)(const void *ctx,
 					    const struct bpf_insn *insn);
 	struct bpf_prog_aux	*aux;		/* Auxiliary fields */
 	struct sock_fprog_kern	*orig_prog;	/* Original BPF program */
+	void __percpu		*private_stack_ptr;
 	/* Instructions for interpreter */
 	union {
 		DECLARE_FLEX_ARRAY(struct sock_filter, insns);
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index 0727fff6de0e..d6eb052f6631 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -1239,6 +1239,7 @@  void __weak bpf_jit_free(struct bpf_prog *fp)
 		struct bpf_binary_header *hdr = bpf_jit_binary_hdr(fp);
 
 		bpf_jit_binary_free(hdr);
+		free_percpu(fp->private_stack_ptr);
 		WARN_ON_ONCE(!bpf_prog_kallsyms_verify_off(fp));
 	}
 
@@ -2420,6 +2421,24 @@  struct bpf_prog *bpf_prog_select_runtime(struct bpf_prog *fp, int *err)
 		if (*err)
 			return fp;
 
+		if (fp->pstack_eligible) {
+			if (!fp->aux->stack_depth) {
+				fp->pstack = PSTACK_TREE_NO;
+			} else {
+				void __percpu *private_stack_ptr;
+
+				fp->pstack = PSTACK_TREE_ROOT;
+				private_stack_ptr =
+					__alloc_percpu_gfp(fp->aux->stack_depth, 8, GFP_KERNEL);
+				if (!private_stack_ptr) {
+					*err = -ENOMEM;
+					return fp;
+				}
+				fp->subtree_stack_depth = fp->aux->stack_depth;
+				fp->private_stack_ptr = private_stack_ptr;
+			}
+		}
+
 		fp = bpf_int_jit_compile(fp);
 		bpf_prog_jit_attempt_done(fp);
 		if (!fp->jited && jit_needed) {
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 69e17cb22037..9d093e2013ca 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -20060,6 +20060,7 @@  static int jit_subprogs(struct bpf_verifier_env *env)
 {
 	struct bpf_prog *prog = env->prog, **func, *tmp;
 	int i, j, subprog_start, subprog_end = 0, len, subprog;
+	int subtree_top_idx, subtree_stack_depth;
 	struct bpf_map *map_ptr;
 	struct bpf_insn *insn;
 	void *old_bpf_func;
@@ -20138,6 +20139,35 @@  static int jit_subprogs(struct bpf_verifier_env *env)
 		func[i]->is_func = 1;
 		func[i]->sleepable = prog->sleepable;
 		func[i]->aux->func_idx = i;
+
+		subtree_top_idx = env->subprog_info[i].subtree_top_idx;
+		if (env->subprog_info[subtree_top_idx].pstack_eligible) {
+			if (subtree_top_idx == i)
+				func[i]->subtree_stack_depth =
+					env->subprog_info[i].subtree_stack_depth;
+
+			subtree_stack_depth = func[i]->subtree_stack_depth;
+			if (subtree_top_idx != i) {
+				if (env->subprog_info[subtree_top_idx].subtree_stack_depth)
+					func[i]->pstack = PSTACK_TREE_INTERNAL;
+				else
+					func[i]->pstack = PSTACK_TREE_NO;
+			} else if (!subtree_stack_depth) {
+				func[i]->pstack = PSTACK_TREE_INTERNAL;
+			} else {
+				void __percpu *private_stack_ptr;
+
+				func[i]->pstack = PSTACK_TREE_ROOT;
+				private_stack_ptr =
+					__alloc_percpu_gfp(subtree_stack_depth, 8, GFP_KERNEL);
+				if (!private_stack_ptr) {
+					err = -ENOMEM;
+					goto out_free;
+				}
+				func[i]->private_stack_ptr = private_stack_ptr;
+			}
+		}
+
 		/* Below members will be freed only at prog->aux */
 		func[i]->aux->btf = prog->aux->btf;
 		func[i]->aux->func_info = prog->aux->func_info;