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 | 22 ++++++++++++++++++++++ kernel/xsched/vstream.c | 24 ++++++++++++++++++++++++ 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 ec78604c3042..fae600ef2cac 100644 --- a/kernel/xsched/dmem.c +++ b/kernel/xsched/dmem.c @@ -202,3 +202,25 @@ 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; + + if (!ctx || !xsched_dmem_used()) + return; + + spin_lock(&ctx->ctx_lock); + + list_for_each_entry_safe(pool, tmp, &ctx->pool_list, pool_node) { + list_del(&pool->pool_node); + spin_unlock(&ctx->ctx_lock); + + dmem_cgroup_uncharge(pool->pool, pool->size); + kfree(pool); + + spin_lock(&ctx->ctx_lock); + } + + spin_unlock(&ctx->ctx_lock); +} diff --git a/kernel/xsched/vstream.c b/kernel/xsched/vstream.c index 593a2a3519d0..8ead548f40cb 100644 --- a/kernel/xsched/vstream.c +++ b/kernel/xsched/vstream.c @@ -696,9 +696,32 @@ static int vstream_hbm_free(struct vstream_args *arg) return xsched_dmem_free(ctx, arg); } + +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); + mutex_unlock(&xcu_found->ctx_list_lock); + + if (ctx) + xsched_dmem_cleanup(ctx); + + 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 */ /* @@ -710,6 +733,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