[PATCH RFC v6 3/5] mm/migrate: add copy offload registration infrastructure

[Shivank Garg <shivankg@amd.com>]

Add a registration interface that lets a single offload provider
(DMA, multi-threaded CPU copy, etc) take over the batch folio copy
performed by migrate_pages_batch().

The provider fills in a struct migrator with an offload_copy()
callback and calls migrate_offload_register(), passing the set of
migration reasons it wants to handle. Registration updates the
migrate_offload_batch_copy_fn() static_call and enables the
migrate_offload_enabled static branch. migrate_offload_unregister()
reverts both.

The active_reason_mask is the single source of truth for which migration
reasons are eligible for batched copy and offload. This is checked with
the migrate_should_offload(). The active reason mask can be changed at
runtime via migrate_offload_set_reason_mask(). Also, add
migrate_offload_reason_mask_parse() and
migrate_offload_reason_mask_format() so a provider can expose its reason
mask using migrate_reason_names[] (e.g. "compaction,demotion", "all",
"none").

Only one migrator can be active at a time. A second registration
gets -EBUSY, and only the active migrator can unregister itself.
The static_call dispatch is protected by SRCU so that the
synchronize_srcu() in unregister waits for all in-flight copy before
the module reference is dropped.

Callbacks must set FOLIO_CONTENT_COPIED on each successfully copied
dst folio so __migrate_folio() skips the per-folio copy; folios
without the marker fall back to CPU copy in the move phase.

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>
---
 MAINTAINERS                          |   2 +
 include/linux/migrate_copy_offload.h |  69 ++++++++++
 mm/Kconfig                           |   6 +
 mm/Makefile                          |   1 +
 mm/migrate.c                         |   9 +-
 mm/migrate_copy_offload.c            | 249 +++++++++++++++++++++++++++++++++++
 6 files changed, 329 insertions(+), 7 deletions(-)

diff --git a/MAINTAINERS b/MAINTAINERS
index 762355b62d54..fa32b447f3e5 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -17065,9 +17065,11 @@ F:	Documentation/ABI/testing/sysfs-kernel-mm-mempolicy-weighted-interleave
 F:	include/linux/mempolicy.h
 F:	include/uapi/linux/mempolicy.h
 F:	include/linux/migrate.h
+F:	include/linux/migrate_copy_offload.h
 F:	include/linux/migrate_mode.h
 F:	mm/mempolicy.c
 F:	mm/migrate.c
+F:	mm/migrate_copy_offload.c
 F:	mm/migrate_device.c
 
 MEMORY MANAGEMENT - MGLRU (MULTI-GEN LRU)
diff --git a/include/linux/migrate_copy_offload.h b/include/linux/migrate_copy_offload.h
new file mode 100644
index 000000000000..97941ebd4979
--- /dev/null
+++ b/include/linux/migrate_copy_offload.h
@@ -0,0 +1,69 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef _LINUX_MIGRATE_COPY_OFFLOAD_H
+#define _LINUX_MIGRATE_COPY_OFFLOAD_H
+
+#include <linux/bits.h>
+#include <linux/errno.h>
+#include <linux/migrate_mode.h>
+
+struct list_head;
+struct module;
+
+#define MIGRATOR_NAME_LEN 32
+
+/*
+ * Reasons that may ever be offloaded. A driver's mask is clamped to this.
+ * Allow all reasons by default.
+ */
+#define MIGRATE_OFFLOAD_REASONS_ALLOWED		(BIT(MR_TYPES) - 1)
+
+/**
+ * struct migrator - batch-copy provider for page migration.
+ * @name: name of the provider.
+ * @offload_copy: copy @folio_cnt folios from @src_list to @dst_list.
+ *
+ *	The migrator may inspect @folio_cnt to decide whether the batch
+ *	is worth offloading, e.g. skip when the batch is too small to
+ *	amortize setup cost.
+ *	The callback must set FOLIO_CONTENT_COPIED in dst->migrate_info
+ *	for each successfully copied destination folio.
+ *	Folios without this marker are copied via per-folio CPU copy in
+ *	the move phase.
+ *
+ * @owner: module providing the migrator. NULL for built-in (=y) drivers.
+ */
+struct migrator {
+	char name[MIGRATOR_NAME_LEN];
+	int (*offload_copy)(struct list_head *dst_list, struct list_head *src_list,
+			unsigned int folio_cnt);
+	struct module *owner;
+};
+
+#ifdef CONFIG_MIGRATION_COPY_OFFLOAD
+int migrate_offload_register(const struct migrator *m, unsigned long reason_mask);
+int migrate_offload_unregister(const struct migrator *m);
+int migrate_offload_set_reason_mask(const struct migrator *m, unsigned long mask);
+int migrate_offload_reason_mask_parse(const char *buf, unsigned long *maskp);
+int migrate_offload_reason_mask_format(char *buf, unsigned long mask);
+bool migrate_should_offload(int reason);
+int migrate_offload_batch_copy(struct list_head *dst_batch,
+		struct list_head *src_batch, unsigned int nr_batch);
+#else
+static inline int migrate_offload_register(const struct migrator *m,
+		unsigned long reason_mask) { return -EOPNOTSUPP; }
+static inline int migrate_offload_unregister(const struct migrator *m) { return -EOPNOTSUPP; }
+static inline int migrate_offload_set_reason_mask(const struct migrator *m,
+		unsigned long mask) { return -EOPNOTSUPP; }
+static inline int migrate_offload_reason_mask_parse(const char *buf,
+		unsigned long *maskp) { return -EOPNOTSUPP; }
+static inline int migrate_offload_reason_mask_format(char *buf,
+		unsigned long mask) { return -EOPNOTSUPP; }
+static inline bool migrate_should_offload(int reason) { return false; }
+static inline int migrate_offload_batch_copy(struct list_head *dst_batch,
+		struct list_head *src_batch, unsigned int nr_batch)
+{
+	return -EOPNOTSUPP;
+}
+#endif
+
+#endif /* _LINUX_MIGRATE_COPY_OFFLOAD_H */
diff --git a/mm/Kconfig b/mm/Kconfig
index 9e0ca4824905..c65639bae268 100644
--- a/mm/Kconfig
+++ b/mm/Kconfig
@@ -700,6 +700,12 @@ config MIGRATION
 config DEVICE_MIGRATION
 	def_bool MIGRATION && ZONE_DEVICE
 
+# Page-migration batch-copy offload infrastructure.
+# Selected by migrator drivers (e.g. CONFIG_DCBM_DMA).
+config MIGRATION_COPY_OFFLOAD
+	bool
+	depends on MIGRATION && 64BIT
+
 config ARCH_ENABLE_HUGEPAGE_MIGRATION
 	bool
 
diff --git a/mm/Makefile b/mm/Makefile
index eff9f9e7e061..256fb532b672 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 41b732e78a67..4fed3110ca0a 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -43,6 +43,7 @@
 #include <linux/sched/sysctl.h>
 #include <linux/memory-tiers.h>
 #include <linux/pagewalk.h>
+#include <linux/migrate_copy_offload.h>
 
 #include <asm/tlbflush.h>
 
@@ -51,12 +52,6 @@
 #include "internal.h"
 #include "swap.h"
 
-/* For now, never offload. Wired up in later patch. */
-static bool migrate_should_offload(int reason)
-{
-	return false;
-}
-
 static const struct movable_operations *offline_movable_ops;
 static const struct movable_operations *zsmalloc_movable_ops;
 
@@ -2050,7 +2045,7 @@ static int migrate_pages_batch(struct list_head *from,
 
 	/* Batch-copy eligible folios before the move phase */
 	if (!list_empty(&unmap_batch))
-		migrate_folios_mc_copy(&dst_batch, &unmap_batch, nr_batch);
+		migrate_offload_batch_copy(&dst_batch, &unmap_batch, nr_batch);
 
 	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..335db8e80556
--- /dev/null
+++ b/mm/migrate_copy_offload.c
@@ -0,0 +1,249 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/bitops.h>
+#include <linux/jump_label.h>
+#include <linux/kstrtox.h>
+#include <linux/migrate.h>
+#include <linux/migrate_copy_offload.h>
+#include <linux/mm.h>
+#include <linux/module.h>
+#include <linux/slab.h>
+#include <linux/srcu.h>
+#include <linux/static_call.h>
+#include <linux/string.h>
+#include <linux/sysfs.h>
+
+DEFINE_STATIC_KEY_FALSE(migrate_offload_enabled);
+DEFINE_SRCU(migrate_offload_srcu);
+DEFINE_STATIC_CALL(migrate_offload_batch_copy_fn, migrate_folios_mc_copy);
+
+static DEFINE_MUTEX(migrator_mutex);
+static const struct migrator *active_migrator;
+
+/*
+ * The active migrator's reason mask. This is the single source of truth
+ * for which reasons are offloaded.
+ */
+static unsigned long active_reason_mask;
+
+bool migrate_should_offload(int reason)
+{
+	if (!static_branch_unlikely(&migrate_offload_enabled))
+		return false;
+
+	return READ_ONCE(active_reason_mask) & BIT(reason);
+}
+
+/**
+ * migrate_offload_reason_mask_parse - parse migration reason mask.
+ * @buf: input string, either a number (hex/decimal, e.g. "0x101") or a
+ *	comma-separated list of reason names (see migrate_reason_names[]),
+ *	e.g. "compaction,demotion". The tokens "all" and "none" select
+ *	every or no reason.
+ * @maskp: parsed mask, clamped to MIGRATE_OFFLOAD_REASONS_ALLOWED.
+ *
+ * Return: 0 on success, -EINVAL on an unknown name, -ENOMEM on OOM.
+ */
+int migrate_offload_reason_mask_parse(const char *buf, unsigned long *maskp)
+{
+	unsigned long mask;
+	char *copy, *p, *tok;
+	int ret = 0;
+
+	/* A plain number is treated as a raw mask. */
+	if (!kstrtoul(buf, 0, &mask))
+		goto done;
+
+	copy = kstrdup(buf, GFP_KERNEL);
+	if (!copy)
+		return -ENOMEM;
+	mask = 0;
+	p = strim(copy);
+
+	while ((tok = strsep(&p, ",")) != NULL) {
+		int i;
+
+		tok = strim(tok);
+		if (!*tok)
+			continue;
+		if (!strcmp(tok, "none")) {
+			mask = 0;
+			continue;
+		}
+		if (!strcmp(tok, "all")) {
+			mask = MIGRATE_OFFLOAD_REASONS_ALLOWED;
+			continue;
+		}
+		i = match_string(migrate_reason_names, MR_TYPES, tok);
+		if (i < 0) {
+			ret = -EINVAL;
+			break;
+		}
+		mask |= BIT(i);
+	}
+	kfree(copy);
+	if (ret)
+		return ret;
+done:
+	*maskp = mask & MIGRATE_OFFLOAD_REASONS_ALLOWED;
+	return 0;
+}
+EXPORT_SYMBOL_GPL(migrate_offload_reason_mask_parse);
+
+/**
+ * migrate_offload_reason_mask_format - render a reason mask as names.
+ * @buf: output buffer from kernel_param_ops get.
+ * @mask: reason mask to render.
+ *
+ * Emits a comma-separated list of reason names, or "none".
+ *
+ * Return: number of bytes written to @buf.
+ */
+int migrate_offload_reason_mask_format(char *buf, unsigned long mask)
+{
+	int len = 0, i;
+
+	for_each_set_bit(i, &mask, MR_TYPES)
+		len += sysfs_emit_at(buf, len, "%s%s",
+				     len ? "," : "", migrate_reason_names[i]);
+	if (!len)
+		return sysfs_emit(buf, "none\n");
+
+	len += sysfs_emit_at(buf, len, "\n");
+	return len;
+}
+EXPORT_SYMBOL_GPL(migrate_offload_reason_mask_format);
+
+/*
+ * Hand the batch to the registered migrator. The migrator may decline
+ * (typically based on batch size), in which case the move phase falls
+ * back to per-folio CPU copy.
+ */
+int migrate_offload_batch_copy(struct list_head *dst_batch,
+		struct list_head *src_batch, unsigned int nr_batch)
+{
+	int idx, rc;
+
+	idx = srcu_read_lock(&migrate_offload_srcu);
+	rc = static_call(migrate_offload_batch_copy_fn)(dst_batch, src_batch, nr_batch);
+	srcu_read_unlock(&migrate_offload_srcu, idx);
+	return rc;
+}
+
+/**
+ * migrate_offload_set_reason_mask - update the active migrator's reason mask.
+ * @m: migrator (must be the currently active one).
+ * @mask: new reason mask.
+ *
+ * Return: 0 on success, -EINVAL if @m is NULL or not the active migrator.
+ */
+int migrate_offload_set_reason_mask(const struct migrator *m, unsigned long mask)
+{
+	if (!m)
+		return -EINVAL;
+
+	mask &= MIGRATE_OFFLOAD_REASONS_ALLOWED;
+
+	mutex_lock(&migrator_mutex);
+	if (active_migrator != m) {
+		mutex_unlock(&migrator_mutex);
+		return -EINVAL;
+	}
+	WRITE_ONCE(active_reason_mask, mask);
+	mutex_unlock(&migrator_mutex);
+	return 0;
+}
+EXPORT_SYMBOL_GPL(migrate_offload_set_reason_mask);
+
+/**
+ * migrate_offload_register - register a batch-copy provider for page migration.
+ * @m: migrator to install.
+ * @reason_mask: initial set of BIT(MR_*) reasons to offload, can be changed
+ *	later via migrate_offload_set_reason_mask().
+ *
+ * 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_register(const struct migrator *m, unsigned long reason_mask)
+{
+	unsigned long mask;
+	int ret = 0;
+
+	if (!m || !m->offload_copy)
+		return -EINVAL;
+
+	mask = reason_mask & MIGRATE_OFFLOAD_REASONS_ALLOWED;
+
+	mutex_lock(&migrator_mutex);
+	if (active_migrator) {
+		ret = -EBUSY;
+		goto unlock;
+	}
+
+	/* @owner is NULL for built-in (=y) drivers; nothing to ref. */
+	if (m->owner && !try_module_get(m->owner)) {
+		ret = -ENODEV;
+		goto unlock;
+	}
+
+	WRITE_ONCE(active_reason_mask, mask);
+	static_call_update(migrate_offload_batch_copy_fn, m->offload_copy);
+	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 (reason_mask=0x%lx)\n",
+			m->name, mask);
+	return ret;
+}
+EXPORT_SYMBOL_GPL(migrate_offload_register);
+
+/**
+ * migrate_offload_unregister - 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_unregister(const struct migrator *m)
+{
+	struct module *owner;
+
+	if (!m)
+		return -EINVAL;
+
+	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
+	 * cannot enter the batch path.
+	 */
+	static_branch_disable(&migrate_offload_enabled);
+	WRITE_ONCE(active_reason_mask, 0);
+	static_call_update(migrate_offload_batch_copy_fn, migrate_folios_mc_copy);
+	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_unregister);

-- 
2.43.0