Re: [RFC PATCH v4 4/6] mm/migrate: add copy offload registration infrastructure

["Huang, Ying" <ying.huang@linux.alibaba.com>]

Shivank Garg <shivankg@amd.com> writes:

> Introduce CONFIG_MIGRATION_COPY_OFFLOAD, which lets offload driver

Do we really need a new kconfig option?  IMHO, we have too many now.
Because we have a jump label already, the performance difference should
be trivial.  Can you measure the size difference?

> (DMA, multi-threaded CPU copy, etc) take over the batch folio copy in
> migrate_pages_batch().
>
> Offload driver fill in a struct migrator with their offload_copy() and
> should_batch() implementation and call migrate_offload_start(), which
> patches the migrate_offload_copy() static_call and flips the
> migrate_offload_enabled static branch. The migrate_offload_stop() call
> reverts both.
>
> Only one migrator can be active a time. A second registration returns
> -EBUSY, and only the active migrator can stop itself. The static_call
> dispatch is under SRCU so synchronize_srcu() in stop path guarantees
> no in-flight copy before the module reference is dropped.
>
> Co-developed-by: Mike Day <michael.day@amd.com>
> Signed-off-by: Mike Day <michael.day@amd.com>
> Signed-off-by: Shivank Garg <shivankg@amd.com>
> ---
>  include/linux/migrate_copy_offload.h | 34 ++++++++++
>  mm/Kconfig                           |  9 +++
>  mm/Makefile                          |  1 +
>  mm/migrate.c                         | 30 ++++++++-
>  mm/migrate_copy_offload.c            | 99 ++++++++++++++++++++++++++++
>  5 files changed, 171 insertions(+), 2 deletions(-)
>  create mode 100644 include/linux/migrate_copy_offload.h
>  create mode 100644 mm/migrate_copy_offload.c
>
> diff --git a/include/linux/migrate_copy_offload.h b/include/linux/migrate_copy_offload.h
> new file mode 100644
> index 000000000000..ee112826ebdf
> --- /dev/null
> +++ b/include/linux/migrate_copy_offload.h
> @@ -0,0 +1,34 @@
> +/* SPDX-License-Identifier: GPL-2.0 */
> +#ifndef _LINUX_MIGRATE_COPY_OFFLOAD_H
> +#define _LINUX_MIGRATE_COPY_OFFLOAD_H
> +
> +#include <linux/jump_label.h>
> +#include <linux/srcu.h>
> +#include <linux/types.h>
> +
> +struct list_head;
> +struct module;
> +
> +#define MIGRATOR_NAME_LEN 32
> +
> +struct migrator {
> +	char name[MIGRATOR_NAME_LEN];
> +	int (*offload_copy)(struct list_head *dst_list,
> +			    struct list_head *src_list,
> +			    unsigned int folio_cnt);
> +	bool (*should_batch)(int reason);
> +	struct module *owner;
> +};
> +
> +#ifdef CONFIG_MIGRATION_COPY_OFFLOAD
> +extern struct static_key_false migrate_offload_enabled;
> +extern struct srcu_struct migrate_offload_srcu;
> +bool migrate_should_batch_default(int reason);
> +int migrate_offload_start(struct migrator *m);
> +int migrate_offload_stop(struct migrator *m);

Why not naming the function migrate_offload_register/unregister()?
IMHO, that sounds more natural.

> +#else
> +static inline int migrate_offload_start(struct migrator *m) { return 0; }
> +static inline int migrate_offload_stop(struct migrator *m) { return 0; }
> +#endif /* CONFIG_MIGRATION_COPY_OFFLOAD */
> +
> +#endif /* _LINUX_MIGRATE_COPY_OFFLOAD_H */
> diff --git a/mm/Kconfig b/mm/Kconfig
> index ebd8ea353687..faf0cae9991b 100644
> --- a/mm/Kconfig
> +++ b/mm/Kconfig
> @@ -648,6 +648,15 @@ config MIGRATION
>  config DEVICE_MIGRATION
>  	def_bool MIGRATION && ZONE_DEVICE
>  
> +config MIGRATION_COPY_OFFLOAD
> +	bool "Page migration copy offload"
> +	depends on MIGRATION
> +	help
> +	  Adds migration copy offload infrastructure which allow
> +	  offload engines (DMA, multi-threaded CPU copy, etc.) to
> +	  register as the batch-copy provider for page migration
> +	  via migrate_offload_start()/migrate_offload_stop().
> +
>  config ARCH_ENABLE_HUGEPAGE_MIGRATION
>  	bool
>  
> diff --git a/mm/Makefile b/mm/Makefile
> index 8ad2ab08244e..db1ac8097089 100644
> --- a/mm/Makefile
> +++ b/mm/Makefile
> @@ -96,6 +96,7 @@ obj-$(CONFIG_FAILSLAB) += failslab.o
>  obj-$(CONFIG_FAIL_PAGE_ALLOC) += fail_page_alloc.o
>  obj-$(CONFIG_MEMTEST)		+= memtest.o
>  obj-$(CONFIG_MIGRATION) += migrate.o
> +obj-$(CONFIG_MIGRATION_COPY_OFFLOAD) += migrate_copy_offload.o
>  obj-$(CONFIG_NUMA) += memory-tiers.o
>  obj-$(CONFIG_DEVICE_MIGRATION) += migrate_device.o
>  obj-$(CONFIG_TRANSPARENT_HUGEPAGE) += huge_memory.o khugepaged.o
> diff --git a/mm/migrate.c b/mm/migrate.c
> index 69daa16f9cf3..acaaa9cc0d4f 100644
> --- a/mm/migrate.c
> +++ b/mm/migrate.c
> @@ -44,6 +44,8 @@
>  #include <linux/memory-tiers.h>
>  #include <linux/pagewalk.h>
>  #include <linux/jump_label.h>
> +#include <linux/static_call.h>
> +#include <linux/migrate_copy_offload.h>
>  
>  #include <asm/tlbflush.h>
>  
> @@ -54,6 +56,17 @@
>  
>  DEFINE_STATIC_KEY_FALSE(migrate_offload_enabled);
>  
> +#ifdef CONFIG_MIGRATION_COPY_OFFLOAD
> +DEFINE_SRCU(migrate_offload_srcu);
> +DEFINE_STATIC_CALL(migrate_offload_copy, folios_mc_copy);
> +
> +bool migrate_should_batch_default(int reason)
> +{
> +	return false;
> +}
> +DEFINE_STATIC_CALL(migrate_should_batch, migrate_should_batch_default);
> +#endif
> +
>  static const struct movable_operations *offline_movable_ops;
>  static const struct movable_operations *zsmalloc_movable_ops;
>  
> @@ -1820,11 +1833,18 @@ static int migrate_pages_batch(struct list_head *from,
>  	LIST_HEAD(dst_batch);
>  	LIST_HEAD(src_std);
>  	LIST_HEAD(dst_std);
> +	bool do_batch = false;
>  	bool nosplit = (reason == MR_NUMA_MISPLACED);
>  
>  	VM_WARN_ON_ONCE(mode != MIGRATE_ASYNC &&
>  			!list_empty(from) && !list_is_singular(from));
>  
> +#ifdef CONFIG_MIGRATION_COPY_OFFLOAD
> +	/* Check if the offload driver wants to batch for this reason */
> +	if (static_branch_unlikely(&migrate_offload_enabled))
> +		do_batch = static_call(migrate_should_batch)(reason);

Should batching based on "reason" be determined by the general migrate
code instead of the migrator implementation?  For example, if we only
batch copying for ASYNC migration, we should determine that in
migrate_pages_batch() instead of the migreation implementation.  Or am I
missed something?  If so, can you provide an example?

> +#endif
> +
>  	for (pass = 0; pass < nr_pass && retry; pass++) {
>  		retry = 0;
>  		thp_retry = 0;
> @@ -1967,7 +1987,7 @@ static int migrate_pages_batch(struct list_head *from,
>  				break;
>  			case 0:
>  				if (static_branch_unlikely(&migrate_offload_enabled) &&
> -				    folio_supports_batch_copy(folio)) {
> +				    do_batch && folio_supports_batch_copy(folio)) {
>  					list_move_tail(&folio->lru, &src_batch);
>  					list_add_tail(&dst->lru, &dst_batch);
>  					nr_batch++;
> @@ -1997,11 +2017,17 @@ static int migrate_pages_batch(struct list_head *from,
>  	/* Flush TLBs for all unmapped folios */
>  	try_to_unmap_flush();
>  
> +#ifdef CONFIG_MIGRATION_COPY_OFFLOAD
>  	/* Batch-copy eligible folios before the move phase */
>  	if (!list_empty(&src_batch)) {

Guard with "static_branch_unlikely(&migrate_offload_enabled)" first?
Better to define a inline function to shorten the expression.

> -		rc = folios_mc_copy(&dst_batch, &src_batch, nr_batch);
> +		int idx = srcu_read_lock(&migrate_offload_srcu);
> +
> +		rc = static_call(migrate_offload_copy)(&dst_batch,
> +				&src_batch, nr_batch);
> +		srcu_read_unlock(&migrate_offload_srcu, idx);
>  		batch_copied = (rc == 0);
>  	}
> +#endif
>  
>  	retry = 1;
>  	for (pass = 0; pass < nr_pass && retry; pass++) {
> diff --git a/mm/migrate_copy_offload.c b/mm/migrate_copy_offload.c
> new file mode 100644
> index 000000000000..c22068fe09a0
> --- /dev/null
> +++ b/mm/migrate_copy_offload.c
> @@ -0,0 +1,99 @@
> +// SPDX-License-Identifier: GPL-2.0
> +#include <linux/jump_label.h>
> +#include <linux/module.h>
> +#include <linux/srcu.h>
> +#include <linux/migrate.h>
> +#include <linux/migrate_copy_offload.h>
> +#include <linux/static_call.h>
> +
> +static DEFINE_MUTEX(migrator_mutex);
> +static struct migrator *active_migrator;
> +
> +DECLARE_STATIC_CALL(migrate_offload_copy, folios_mc_copy);
> +DECLARE_STATIC_CALL(migrate_should_batch, migrate_should_batch_default);
> +
> +/**
> + * migrate_offload_start - register a batch-copy provider for page migration.
> + * @m: migrator to install.
> + *
> + * Only one provider can be active at a time, returns -EBUSY if another migrator
> + * is already registered.
> + *
> + * Return: 0 on success, negative errno on failure.
> + */
> +int migrate_offload_start(struct migrator *m)
> +{
> +	int ret = 0;
> +
> +	if (!m || !m->offload_copy)
> +		return -EINVAL;
> +
> +	mutex_lock(&migrator_mutex);
> +	if (active_migrator) {
> +		ret = -EBUSY;
> +		goto unlock;
> +	}
> +
> +	if (m->owner && !try_module_get(m->owner)) {
> +		ret = -ENODEV;
> +		goto unlock;
> +	}
> +
> +	static_call_update(migrate_offload_copy, m->offload_copy);
> +	static_call_update(migrate_should_batch,
> +		m->should_batch ? m->should_batch : migrate_should_batch_default);
> +	active_migrator = m;
> +	static_branch_enable(&migrate_offload_enabled);
> +
> +unlock:
> +	mutex_unlock(&migrator_mutex);
> +
> +	if (ret)
> +		pr_err("migrate_offload: %s: failed to register (%d)\n",
> +		       m->name, ret);
> +	else
> +		pr_info("migrate_offload: enabled by %s\n", m->name);
> +	return ret;
> +}
> +EXPORT_SYMBOL_GPL(migrate_offload_start);
> +
> +/**
> + * migrate_offload_stop - unregister the active batch-copy provider.
> + * @m: migrator to remove (must be the currently active one).
> + *
> + * Reverts static_call targets and waits for SRCU grace period so that
> + * no in-flight migration is still calling the driver functions before
> + * releasing the module.
> + *
> + * Return: 0 on success, negative errno on failure.
> + */
> +int migrate_offload_stop(struct migrator *m)
> +{
> +	struct module *owner;
> +
> +	mutex_lock(&migrator_mutex);
> +	if (active_migrator != m) {
> +		mutex_unlock(&migrator_mutex);
> +		return -EINVAL;
> +	}
> +
> +	/*
> +	 * Disable the static branch first so new migrate_pages_batch calls
> +	 * won't enter the batch copy path.
> +	 */
> +	static_branch_disable(&migrate_offload_enabled);
> +	static_call_update(migrate_offload_copy, folios_mc_copy);
> +	static_call_update(migrate_should_batch, migrate_should_batch_default);
> +	owner = active_migrator->owner;
> +	active_migrator = NULL;
> +	mutex_unlock(&migrator_mutex);
> +
> +	/* Wait for all in-flight callers to finish before module_put(). */
> +	synchronize_srcu(&migrate_offload_srcu);
> +	if (owner)
> +		module_put(owner);
> +
> +	pr_info("migrate_offload: disabled by %s\n", m->name);
> +	return 0;
> +}
> +EXPORT_SYMBOL_GPL(migrate_offload_stop);

---
Best Regards,
Huang, Ying