bpf:programmable: Add helper func to set cpus_ptr in task

hulk inclusion
category: feature
bugzilla: https://gitee.com/openeuler/kernel/issues/I5KUFB
CVE: NA

--------------------------------

Add helper funciton to set cpus_ptr in task.
Signed-off-by: NHui Tang <tanghui20@huawei.com>

include/uapi/linux/bpf.h

@@ -3906,6 +3906,12 @@ union bpf_attr {
  *		check src_cpu whether share cache with dst_cpu.
  *	Return
  *		yes 1, no 0.
+ *
+ * int bpf_sched_set_task_cpus_ptr(struct sched_migrate_ctx *h_ctx, struct cpumask *cpus, int len)
+ *	Description
+ *		set cpus_ptr in task.
+ *	Return
+ *		0 on success, or a negative error in case of failure.
  */
 #define __BPF_FUNC_MAPPER(FN)		\
 	FN(unspec),			\
@@ -4084,6 +4090,7 @@ union bpf_attr {
 	FN(get_cpumask_info),		\
 	FN(cpumask_op),			\
 	FN(cpus_share_cache),		\
+	FN(sched_set_task_cpus_ptr),	\
 	/* */
 
 /* integer value in 'imm' field of BPF_CALL instruction selects which helper

kernel/sched/bpf_sched.c

@@ -503,6 +503,27 @@ static const struct bpf_func_proto bpf_cpus_share_cache_proto = {
 	.arg2_type	= ARG_ANYTHING,
 };
 
+BPF_CALL_3(bpf_sched_set_task_cpus_ptr, struct sched_migrate_ctx *, h_ctx,
+	   struct cpumask *, cpus, int, len)
+{
+	if (len != sizeof(*cpus))
+		return -EINVAL;
+
+	h_ctx->task->cpus_ptr = cpus;
+	return 0;
+}
+
+BTF_ID_LIST_SINGLE(bpf_sched_migrate_ctx_ids, struct, sched_migrate_ctx)
+
+static const struct bpf_func_proto bpf_sched_set_task_cpus_ptr_proto = {
+	.func		= bpf_sched_set_task_cpus_ptr,
+	.gpl_only	= false,
+	.ret_type	= RET_INTEGER,
+	.arg1_type	= ARG_PTR_TO_BTF_ID,
+	.arg1_btf_id	= &bpf_sched_migrate_ctx_ids[0],
+	.arg2_type	= ARG_ANYTHING,
+};
+
 static const struct bpf_func_proto *
 bpf_sched_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 {
@@ -537,6 +558,8 @@ bpf_sched_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 		return &bpf_cpumask_op_proto;
 	case BPF_FUNC_cpus_share_cache:
 		return &bpf_cpus_share_cache_proto;
+	case BPF_FUNC_sched_set_task_cpus_ptr:
+		return &bpf_sched_set_task_cpus_ptr_proto;
 	default:
 		return bpf_base_func_proto(func_id);
 	}

tools/include/uapi/linux/bpf.h

@@ -3906,6 +3906,12 @@ union bpf_attr {
  *		check src_cpu whether share cache with dst_cpu.
  *	Return
  *		true yes, false no.
+ *
+ * int bpf_sched_set_task_cpus_ptr(struct sched_migrate_ctx *h_ctx, struct cpumask *cpus, int len)
+ *	Description
+ *		set cpus_ptr in task.
+ *	Return
+ *		0 on success, or a negative error in case of failure.
  */
 #define __BPF_FUNC_MAPPER(FN)		\
 	FN(unspec),			\
@@ -4084,6 +4090,7 @@ union bpf_attr {
 	FN(get_cpumask_info),		\
 	FN(cpumask_op),			\
 	FN(cpus_share_cache),		\
+	FN(sched_set_task_cpus_ptr),	\
 	/* */
 
 /* integer value in 'imm' field of BPF_CALL instruction selects which helper

tools/lib/bpf/libbpf_sched.h

@@ -463,4 +463,11 @@ static __always_inline  int libbpf_sched_se_tag_of(struct sched_entity *se)
 
 	return se_tag;
 }
+
+static __always_inline void libbpf_sched_set_task_cpus_ptr(
+	struct sched_migrate_ctx *h_ctx,
+	struct cpumask *cpus)
+{
+	bpf_sched_set_task_cpus_ptr(h_ctx, cpus, sizeof(*cpus));
+}
 #endif