From: Alexander Pavlenko <pavlenko.alexander@huawei.com> hulk inclusion category: feature bugzilla: https://atomgit.com/openeuler/kernel/issues/8424 ---------------------------------------- This commit changes the XPU device memory allocation flow: the xsched kernel must first successfully register the intended device memory (dmem) region before any physical memory is allocated from the XPU. Signed-off-by: Alexander Pavlenko <pavlenko.alexander@huawei.com> Signed-off-by: Liu Kai <liukai284@huawei.com> --- include/linux/cgroup_dmem.h | 2 - include/linux/xsched.h | 7 ++- include/uapi/linux/xsched/xcu_vstream.h | 7 +++ kernel/xsched/dmem.c | 73 +++++++++++++++++++++++++ kernel/xsched/vstream.c | 54 ++++++++++++++++++ 5 files changed, 140 insertions(+), 3 deletions(-) diff --git a/include/linux/cgroup_dmem.h b/include/linux/cgroup_dmem.h index b86ca6012516..5730abf5050e 100644 --- a/include/linux/cgroup_dmem.h +++ b/include/linux/cgroup_dmem.h @@ -64,8 +64,6 @@ static inline void dmem_cgroup_pool_state_put(struct dmem_cgroup_pool_state *poo #endif -#ifdef CONFIG_CGROUP_DEVICE bool dmem_cgroup_enabled(void); -#endif #endif /* _CGROUP_DMEM_H */ diff --git a/include/linux/xsched.h b/include/linux/xsched.h index e83b977c4265..d635ecc516c7 100644 --- a/include/linux/xsched.h +++ b/include/linux/xsched.h @@ -374,6 +374,10 @@ struct xsched_context { struct list_head vstream_list; struct list_head ctx_node; +#ifdef CONFIG_CGROUP_DMEM + struct list_head pool_list; +#endif + struct xsched_entity xse; spinlock_t ctx_lock; @@ -515,6 +519,7 @@ static inline u64 xs_calc_delta_fair(u64 delta_exec, u32 weight) int xsched_dmem_register_region(uint64_t size, int dev_id); void xsched_dmem_unregister_region(int dev_id); bool xsched_dmem_used(void); -#endif /* CONFIG_CGROUP_DMEM */ +int xsched_dmem_alloc(struct xsched_context *ctx, struct vstream_args *args); +#endif #endif /* !__LINUX_XSCHED_H__ */ diff --git a/include/uapi/linux/xsched/xcu_vstream.h b/include/uapi/linux/xsched/xcu_vstream.h index 4b3abf386a1c..d4cfa3cf0591 100644 --- a/include/uapi/linux/xsched/xcu_vstream.h +++ b/include/uapi/linux/xsched/xcu_vstream.h @@ -22,6 +22,7 @@ typedef enum VSTREAM_COMMAND { VSTREAM_ALLOC = 0, VSTREAM_FREE, VSTREAM_KICK, + VSTREAM_HBM_ALLOC, MAX_COMMAND } vstream_command_t; @@ -51,6 +52,11 @@ typedef struct vstream_kick_args { KABI_RESERVE_BYTES(2, 8); } vstream_kick_args_t; +typedef struct vstream_hbm_args { + __u64 size; + __u64 addr; +} vstream_hbm_args_t; + typedef struct vstream_args { __u32 channel_id; __u32 fd; @@ -64,6 +70,7 @@ typedef struct vstream_args { vstream_alloc_args_t va_args; vstream_free_args_t vf_args; vstream_kick_args_t vk_args; + vstream_hbm_args_t vm_args; }; __u32 payload_size; diff --git a/kernel/xsched/dmem.c b/kernel/xsched/dmem.c index 1296071740a9..af2e8e51927b 100644 --- a/kernel/xsched/dmem.c +++ b/kernel/xsched/dmem.c @@ -50,6 +50,11 @@ static void xsched_dmem_region_release(struct kref *kref) kfree(region); } +static void xsched_dmem_region_get(struct xsched_dmem_region *region) +{ + kref_get(®ion->kref); +} + static void xsched_dmem_region_put(struct xsched_dmem_region *region) { kref_put(®ion->kref, xsched_dmem_region_release); @@ -71,6 +76,24 @@ struct xsched_dmem_region *find_dmem_region_by_devid(int dev_id) return found; } +static struct xsched_dmem_region *get_dmem_region_by_devid(int dev_id) +{ + struct xsched_dmem_region *region; + + spin_lock(&xsched_dmem_lock); + region = find_dmem_region_by_devid(dev_id); + if (region) + xsched_dmem_region_get(region); + spin_unlock(&xsched_dmem_lock); + + return region; +} + +bool xsched_dmem_used(void) +{ + return READ_ONCE(xsched_dmem_region_num) > 0; +} + int xsched_dmem_register_region(uint64_t size, int dev_id) { struct xsched_dmem_region *region, *found; @@ -141,3 +164,53 @@ void xsched_dmem_unregister_region(int dev_id) XSCHED_INFO("unregister HBM%d region(s) in dmem\n", dev_id); } EXPORT_SYMBOL(xsched_dmem_unregister_region); + +int xsched_dmem_alloc(struct xsched_context *ctx, struct vstream_args *args) +{ + struct dmem_cgroup_pool_state *ret_pool, *ret_limit_pool; + struct xsched_dmem_pool *new_pool; + struct xsched_dmem_region *found; + int ret; + + if (!xsched_dmem_used()) + return -EPERM; + + found = get_dmem_region_by_devid(args->dev_id); + if (!found) { + XSCHED_ERR("Try to charge memory when region is not registered (region HBM%u)\n", + args->dev_id); + return -ENOENT; + } + + ret = dmem_cgroup_try_charge(found->cgroup_region, args->vm_args.size, + &ret_pool, &ret_limit_pool); + if (ret != 0) { + XSCHED_ERR("Fail to charge a new allocation to a HBM region\n"); + dmem_cgroup_pool_state_put(ret_limit_pool); + xsched_dmem_region_put(found); + return ret; + } + + new_pool = kzalloc(sizeof(*new_pool), GFP_KERNEL); + if (!new_pool) { + XSCHED_ERR("Fail to alloc xsched dmem alloc @ %s\n", __func__); + dmem_cgroup_uncharge(ret_pool, args->vm_args.size); + xsched_dmem_region_put(found); + return -ENOMEM; + } + + new_pool->pool = ret_pool; + new_pool->addr = args->vm_args.addr; + new_pool->size = args->vm_args.size; + + spin_lock(&ctx->ctx_lock); + list_add_tail(&new_pool->pool_node, &ctx->pool_list); + spin_unlock(&ctx->ctx_lock); + + xsched_dmem_region_put(found); + + XSCHED_DEBUG("charged %llu bytes, new_alloc = %p with addr %llu", + new_pool->size, new_pool, new_pool->addr); + + return 0; +} diff --git a/kernel/xsched/vstream.c b/kernel/xsched/vstream.c index b39e97682dfb..8f310121f36e 100644 --- a/kernel/xsched/vstream.c +++ b/kernel/xsched/vstream.c @@ -22,6 +22,10 @@ #include <linux/vstream.h> #include <linux/xsched.h> +#ifdef CONFIG_CGROUP_DMEM +#include <linux/cgroup_dmem.h> +#endif + #if defined(CONFIG_XCU_SCHEDULER) && defined(CONFIG_XCU_VSTREAM) #define XCU_HASH_ORDER 6 @@ -169,6 +173,10 @@ static void init_xsched_ctx(struct xsched_context *ctx, INIT_LIST_HEAD(&ctx->vstream_list); INIT_LIST_HEAD(&ctx->ctx_node); +#ifdef CONFIG_CGROUP_DMEM + INIT_LIST_HEAD(&ctx->pool_list); +#endif + spin_lock_init(&ctx->ctx_lock); mutex_init(&ctx->ctx_mutex); } @@ -628,6 +636,51 @@ int vstream_kick(struct vstream_args *arg) return err; } +#ifdef CONFIG_CGROUP_DMEM +static int vstream_hbm_alloc(struct vstream_args *arg) +{ + vstream_info_t vstream_info; + struct xsched_cu *xcu_found; + struct xsched_context *ctx; + int ret = 0; + + if (!dmem_cgroup_enabled()) + return -EPERM; + + xcu_found = xcu_find(XCU_TYPE_XPU, arg->dev_id, arg->channel_id); + if (!xcu_found) + return -EINVAL; + + /* it will either allocate or find a context */ + mutex_lock(&xcu_found->ctx_list_lock); + ctx = ctx_find_by_tgid_and_xcu(current->tgid, xcu_found); + if (ctx) { + kref_get(&ctx->kref); + } else { + vstream_info.tgid = current->tgid; + vstream_info.xcu = xcu_found; + vstream_info.dev_id = arg->dev_id; + vstream_info.channel_id = arg->channel_id; + vstream_info.fd = arg->fd; + + ret = alloc_ctx_from_vstream(&vstream_info, &ctx); + } + mutex_unlock(&xcu_found->ctx_list_lock); + + if (ret != 0) { + XSCHED_ERR("Failed to find a context for HBM alloc"); + return ret; + } + + ret = xsched_dmem_alloc(ctx, arg); + kref_put(&ctx->kref, xsched_task_free); + + return ret; +} +#else +static int vstream_hbm_alloc(struct vstream_args *arg) { return -EOPNOTSUPP; } +#endif /* CONFIG_CGROUP_DMEM */ + /* * vstream_manage_cmd table */ @@ -635,6 +688,7 @@ static vstream_manage_t(*vstream_command_table[MAX_COMMAND + 1]) = { vstream_alloc, // VSTREAM_ALLOC vstream_free, // VSTREAM_FREE vstream_kick, // VSTREAM_KICK + vstream_hbm_alloc, // VSTREAM_HBM_ALLOC NULL // MAX_COMMAND }; -- 2.34.1

hulk inclusion category: feature bugzilla: https://atomgit.com/openeuler/kernel/issues/8424 ---------------------------------------- Introduce the xsched_dmem_cleanup() function to clean up all allocated dmem regions when an AI task completes. This function iterates through and releases all dmem regions allocated by the task, ensuring that device memory resources are properly returned to the system and preventing memory leaks. Signed-off-by: Liu Kai <liukai284@huawei.com> --- include/linux/xsched.h | 1 + include/uapi/linux/xsched/xcu_vstream.h | 1 + kernel/xsched/core.c | 3 +++ kernel/xsched/dmem.c | 18 ++++++++++++++++ kernel/xsched/vstream.c | 28 +++++++++++++++++++++++++ 5 files changed, 51 insertions(+) diff --git a/include/linux/xsched.h b/include/linux/xsched.h index bb29d9243e97..9c6c9fb5dbc4 100644 --- a/include/linux/xsched.h +++ b/include/linux/xsched.h @@ -521,6 +521,7 @@ void xsched_dmem_unregister_region(int dev_id); bool xsched_dmem_used(void); int xsched_dmem_alloc(struct xsched_context *ctx, struct vstream_args *args); int xsched_dmem_free(struct xsched_context *ctx, struct vstream_args *args); +void xsched_dmem_cleanup(struct xsched_context *ctx); #endif #endif /* !__LINUX_XSCHED_H__ */ diff --git a/include/uapi/linux/xsched/xcu_vstream.h b/include/uapi/linux/xsched/xcu_vstream.h index 54c5efb63d53..5e197558e496 100644 --- a/include/uapi/linux/xsched/xcu_vstream.h +++ b/include/uapi/linux/xsched/xcu_vstream.h @@ -24,6 +24,7 @@ typedef enum VSTREAM_COMMAND { VSTREAM_KICK, VSTREAM_HBM_ALLOC, VSTREAM_HBM_FREE, + VSTREAM_HBM_CLEANUP, MAX_COMMAND } vstream_command_t; diff --git a/kernel/xsched/core.c b/kernel/xsched/core.c index 1d835fc44d0c..a4328c09544d 100644 --- a/kernel/xsched/core.c +++ b/kernel/xsched/core.c @@ -174,6 +174,9 @@ int delete_ctx(struct xsched_context *ctx) mutex_unlock(&xcu->xcu_lock); +#ifdef CONFIG_CGROUP_DMEM + xsched_dmem_cleanup(ctx); +#endif xse->class->xse_deinit(xse); return 0; } diff --git a/kernel/xsched/dmem.c b/kernel/xsched/dmem.c index e2e832c45706..8cb3706679b3 100644 --- a/kernel/xsched/dmem.c +++ b/kernel/xsched/dmem.c @@ -243,3 +243,21 @@ int xsched_dmem_free(struct xsched_context *ctx, struct vstream_args *args) return 0; } + +void xsched_dmem_cleanup(struct xsched_context *ctx) +{ + struct xsched_dmem_pool *pool, *tmp; + LIST_HEAD(free_list); + + if (!ctx || !xsched_dmem_used()) + return; + + spin_lock(&ctx->ctx_lock); + list_splice_init(&ctx->pool_list, &free_list); + spin_unlock(&ctx->ctx_lock); + + list_for_each_entry_safe(pool, tmp, &free_list, pool_node) { + dmem_cgroup_uncharge(pool->pool, pool->size); + kfree(pool); + } +} diff --git a/kernel/xsched/vstream.c b/kernel/xsched/vstream.c index a8be0ba8ec39..22974df55186 100644 --- a/kernel/xsched/vstream.c +++ b/kernel/xsched/vstream.c @@ -707,9 +707,36 @@ static int vstream_hbm_free(struct vstream_args *arg) return ret; } + +static int vstream_hbm_cleanup(struct vstream_args *arg) +{ + struct xsched_cu *xcu_found; + struct xsched_context *ctx; + + if (!dmem_cgroup_enabled()) + return -EPERM; + + xcu_found = xcu_find(XCU_TYPE_XPU, arg->dev_id, arg->channel_id); + if (!xcu_found) + return -EINVAL; + + mutex_lock(&xcu_found->ctx_list_lock); + ctx = ctx_find_by_tgid_and_xcu(current->tgid, xcu_found); + if (ctx) + kref_get(&ctx->kref); + mutex_unlock(&xcu_found->ctx_list_lock); + + if (ctx) { + xsched_dmem_cleanup(ctx); + kref_put(&ctx->kref, xsched_task_free); + } + + return 0; +} #else static int vstream_hbm_alloc(struct vstream_args *arg) { return -EOPNOTSUPP; } static int vstream_hbm_free(struct vstream_args *arg) { return -EOPNOTSUPP; } +static int vstream_hbm_cleanup(struct vstream_args *arg) { return -EOPNOTSUPP; } #endif /* CONFIG_CGROUP_DMEM */ /* @@ -721,6 +748,7 @@ static vstream_manage_t(*vstream_command_table[MAX_COMMAND + 1]) = { vstream_kick, // VSTREAM_KICK vstream_hbm_alloc, // VSTREAM_HBM_ALLOC vstream_hbm_free, // VSTREAM_HBM_FREE + vstream_hbm_cleanup, // VSTREAM_HBM_CLEANUP NULL // MAX_COMMAND }; -- 2.34.1