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

Shivank Garg
2026-03-09T12:11:49Z

Introduce CONFIG_MIGRATION_COPY_OFFLOAD, which lets offload driver
(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);
+#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);
+#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)) {
- 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);
--
2.43.0