[PATCH] mm/damon/vaddr: preserve migration weights above UINT_MAX

[Daemyung Kang <charsyam@gmail.com>]

DAMOS migration accepts an unsigned int weight for each destination,
but sums the weights into an unsigned int.  If the total wraps to zero,
no folio is migrated.  Other wrapped totals select the wrong node.

Sum weights into a u64 once per migration walk.  Keep the existing
32-bit divisor path for totals up to UINT_MAX and use a 64-bit remainder
for larger totals.  Add KUnit coverage for the actual summing path and
for destination selection at the overflow boundary.

Fixes: 19c1dc15c859 ("mm/damon/vaddr: use damos->migrate_dests in migrate_{hot,cold}")
Signed-off-by: Daemyung Kang <charsyam@gmail.com>
---
 mm/damon/tests/vaddr-kunit.h | 35 +++++++++++++++++++
 mm/damon/vaddr.c             | 65 +++++++++++++++++++++++-------------
 2 files changed, 77 insertions(+), 23 deletions(-)

diff --git a/mm/damon/tests/vaddr-kunit.h b/mm/damon/tests/vaddr-kunit.h
index d61b503..f67f1a0 100644
--- a/mm/damon/tests/vaddr-kunit.h
+++ b/mm/damon/tests/vaddr-kunit.h
@@ -266,12 +266,47 @@ static void damon_test_apply_three_regions4(struct kunit *test)
 			new_three_regions, expected, ARRAY_SIZE(expected));
 }
 
+static void damon_test_migrate_dests_idx(struct kunit *test)
+{
+	unsigned int weights[] = { UINT_MAX, 1 };
+	struct damos_migrate_dests dests = {
+		.weight_arr = weights,
+		.nr_dests = ARRAY_SIZE(weights),
+	};
+	u64 total = damos_va_migrate_weight_total(&dests);
+
+	KUNIT_EXPECT_EQ(test, total, (u64)UINT_MAX + 1);
+	KUNIT_EXPECT_EQ(test, damos_va_migrate_dests_idx(&dests, 0, total), 0);
+	KUNIT_EXPECT_EQ(test, damos_va_migrate_dests_idx(&dests,
+			UINT_MAX, total), 1);
+
+	weights[1] = 2;
+	total = damos_va_migrate_weight_total(&dests);
+	KUNIT_EXPECT_EQ(test, total, (u64)UINT_MAX + 2);
+	KUNIT_EXPECT_EQ(test, damos_va_migrate_dests_idx(&dests,
+			UINT_MAX, total), 1);
+
+	weights[0] = 0;
+	weights[1] = 0;
+	total = damos_va_migrate_weight_total(&dests);
+	KUNIT_EXPECT_EQ(test, total, 0ULL);
+	KUNIT_EXPECT_EQ(test, damos_va_migrate_dests_idx(&dests, 0, total), -1);
+
+	weights[0] = 1;
+	weights[1] = 3;
+	total = damos_va_migrate_weight_total(&dests);
+	KUNIT_EXPECT_EQ(test, total, 4ULL);
+	KUNIT_EXPECT_EQ(test, damos_va_migrate_dests_idx(&dests, 1, total), 1);
+	KUNIT_EXPECT_EQ(test, damos_va_migrate_dests_idx(&dests, 4, total), 0);
+}
+
 static struct kunit_case damon_test_cases[] = {
 	KUNIT_CASE(damon_test_three_regions_in_vmas),
 	KUNIT_CASE(damon_test_apply_three_regions1),
 	KUNIT_CASE(damon_test_apply_three_regions2),
 	KUNIT_CASE(damon_test_apply_three_regions3),
 	KUNIT_CASE(damon_test_apply_three_regions4),
+	KUNIT_CASE(damon_test_migrate_dests_idx),
 	{},
 };
 
diff --git a/mm/damon/vaddr.c b/mm/damon/vaddr.c
index d5dde97..279dbc6 100644
--- a/mm/damon/vaddr.c
+++ b/mm/damon/vaddr.c
@@ -7,6 +7,7 @@
 
 #include <linux/highmem.h>
 #include <linux/hugetlb.h>
+#include <linux/math64.h>
 #include <linux/mman.h>
 #include <linux/mmu_notifier.h>
 #include <linux/page_idle.h>
@@ -774,8 +775,41 @@ static bool damos_va_filter_out(struct damos *scheme, struct folio *folio,
 struct damos_va_migrate_private {
 	struct list_head *migration_lists;
 	struct damos *scheme;
+	u64 weight_total;
 };
 
+static u64 damos_va_migrate_weight_total(struct damos_migrate_dests *dests)
+{
+	u64 total = 0;
+	size_t i;
+
+	for (i = 0; i < dests->nr_dests; i++)
+		total += dests->weight_arr[i];
+	return total;
+}
+
+static int damos_va_migrate_dests_idx(struct damos_migrate_dests *dests,
+		pgoff_t ilx, u64 weight_total)
+{
+	u64 target;
+	int i;
+
+	if (!weight_total)
+		return -1;
+
+	if (weight_total <= UINT_MAX)
+		target = ilx % (unsigned int)weight_total;
+	else
+		div64_u64_rem(ilx, weight_total, &target);
+
+	for (i = 0; i < dests->nr_dests; i++) {
+		if (target < dests->weight_arr[i])
+			return i;
+		target -= dests->weight_arr[i];
+	}
+	return -1;
+}
+
 /*
  * Place the given folio in the migration_list corresponding to where the folio
  * should be migrated.
@@ -784,13 +818,11 @@ struct damos_va_migrate_private {
  */
 static void damos_va_migrate_dests_add(struct folio *folio,
 		struct vm_area_struct *vma, unsigned long addr,
-		struct damos_migrate_dests *dests,
-		struct list_head *migration_lists)
+		struct damos_va_migrate_private *priv)
 {
+	struct damos_migrate_dests *dests = &priv->scheme->migrate_dests;
 	pgoff_t ilx;
 	int order;
-	unsigned int target;
-	unsigned int weight_total = 0;
 	int i;
 
 	/*
@@ -806,20 +838,10 @@ static void damos_va_migrate_dests_add(struct folio *folio,
 	ilx = vma_start_pgoff(vma) >> order;
 	ilx += linear_page_delta(vma, addr) >> order;
 
-	for (i = 0; i < dests->nr_dests; i++)
-		weight_total += dests->weight_arr[i];
-
-	/* If the total weights are somehow 0, don't migrate at all */
-	if (!weight_total)
+	i = damos_va_migrate_dests_idx(dests, ilx, priv->weight_total);
+	if (i < 0)
 		return;
 
-	target = ilx % weight_total;
-	for (i = 0; i < dests->nr_dests; i++) {
-		if (target < dests->weight_arr[i])
-			break;
-		target -= dests->weight_arr[i];
-	}
-
 	/* If the folio is already in the right node, don't do anything */
 	if (folio_nid(folio) == dests->node_id_arr[i])
 		return;
@@ -829,16 +851,14 @@ static void damos_va_migrate_dests_add(struct folio *folio,
 		return;
 	node_stat_add_folio(folio, NR_ISOLATED_ANON +
 			folio_is_file_lru(folio));
-	list_add(&folio->lru, &migration_lists[i]);
+	list_add(&folio->lru, &priv->migration_lists[i]);
 }
 
 static int damos_va_migrate_pmd_entry(pmd_t *pmd, unsigned long addr,
 		unsigned long next, struct mm_walk *walk)
 {
 	struct damos_va_migrate_private *priv = walk->private;
-	struct list_head *migration_lists = priv->migration_lists;
 	struct damos *s = priv->scheme;
-	struct damos_migrate_dests *dests = &s->migrate_dests;
 	struct folio *folio;
 	spinlock_t *ptl;
 	pte_t *start_pte, *pte, ptent;
@@ -856,8 +876,7 @@ static int damos_va_migrate_pmd_entry(pmd_t *pmd, unsigned long addr,
 			goto huge_out;
 		if (damos_va_filter_out(s, folio, walk->vma, addr, NULL, pmd))
 			goto huge_out;
-		damos_va_migrate_dests_add(folio, walk->vma, addr, dests,
-				migration_lists);
+		damos_va_migrate_dests_add(folio, walk->vma, addr, priv);
 huge_out:
 		spin_unlock(ptl);
 		return 0;
@@ -881,8 +900,7 @@ static int damos_va_migrate_pmd_entry(pmd_t *pmd, unsigned long addr,
 			continue;
 		if (damos_va_filter_out(s, folio, walk->vma, addr, pte, NULL))
 			continue;
-		damos_va_migrate_dests_add(folio, walk->vma, addr, dests,
-				migration_lists);
+		damos_va_migrate_dests_add(folio, walk->vma, addr, priv);
 		page_idx = folio_page_idx(folio, pte_page(ptent));
 		nr = folio_nr_pages(folio) - page_idx;
 	}
@@ -958,6 +976,7 @@ static unsigned long damos_va_migrate(struct damon_target *target,
 	use_target_nid = dests->nr_dests == 0;
 	nr_dests = use_target_nid ? 1 : dests->nr_dests;
 	priv.scheme = s;
+	priv.weight_total = damos_va_migrate_weight_total(dests);
 	priv.migration_lists = kmalloc_objs(*priv.migration_lists, nr_dests);
 	if (!priv.migration_lists)
 		return 0;
-- 
2.43.0