Check that the kmalloc partitions are separate caches by default, and
alias the normal kmalloc caches when they're turned off with
kmalloc_partition=0. The test is skipped on kernels built without
CONFIG_KMALLOC_PARTITION_CACHES.

Make kmalloc_partition_enabled global and export it to KUnit so that
the test can see it.

Signed-off-by: Chris Down <chris@chrisdown.name>
---
 lib/tests/slub_kunit.c | 45 ++++++++++++++++++++++++++++++++++++++++++
 mm/slab.h              |  1 +
 mm/slab_common.c       |  3 ++-
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/lib/tests/slub_kunit.c b/lib/tests/slub_kunit.c
index 1af3f087134f..61227c242ae2 100644
--- a/lib/tests/slub_kunit.c
+++ b/lib/tests/slub_kunit.c
@@ -334,6 +334,50 @@ static void test_slab_buckets(struct kunit *test)
 	kfree(p);
 }
 
+/*
+ * Each partition copy is a separate cache, unless the copies were disabled
+ * with kmalloc_partition=0: then they alias the KMALLOC_NORMAL caches.
+ */
+static void test_kmalloc_partition(struct kunit *test)
+{
+	enum kmalloc_cache_type type;
+	struct kmem_cache *s;
+	unsigned int idx;
+	bool found;
+	void *p;
+
+	if (!IS_ENABLED(CONFIG_KMALLOC_PARTITION_CACHES))
+		kunit_skip(test, "requires CONFIG_KMALLOC_PARTITION_CACHES");
+
+	for (type = KMALLOC_PARTITION_START + 1; type <= KMALLOC_PARTITION_END; type++) {
+		for (idx = 0; idx < ARRAY_SIZE(kmalloc_caches[type]); idx++) {
+			struct kmem_cache *normal = kmalloc_caches[KMALLOC_NORMAL][idx];
+			struct kmem_cache *part = kmalloc_caches[type][idx];
+
+			if (!kmalloc_partition_enabled) {
+				KUNIT_EXPECT_PTR_EQ(test, part, normal);
+			} else if (normal) {
+				KUNIT_EXPECT_NOT_NULL(test, part);
+				KUNIT_EXPECT_PTR_NE(test, part, normal);
+			}
+		}
+	}
+
+	p = kmalloc(64, GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, p);
+	s = virt_to_slab(p)->slab_cache;
+	idx = kmalloc_index(64);
+	if (!kmalloc_partition_enabled) {
+		KUNIT_EXPECT_PTR_EQ(test, s, kmalloc_caches[KMALLOC_NORMAL][idx]);
+	} else {
+		found = false;
+		for (type = KMALLOC_PARTITION_START; type <= KMALLOC_PARTITION_END; type++)
+			found |= s == kmalloc_caches[type][idx];
+		KUNIT_EXPECT_TRUE_MSG(test, found, "allocated from %s", s->name);
+	}
+	kfree(p);
+}
+
 #if defined(CONFIG_PERF_EVENTS) || (defined(CONFIG_KPROBES) && defined(CONFIG_SMP))
 #define NR_ITERATIONS 1000
 #define NR_OBJECTS 1000
@@ -528,6 +572,7 @@ static struct kunit_case test_cases[] = {
 	KUNIT_CASE(test_leak_destroy),
 	KUNIT_CASE(test_krealloc_redzone_zeroing),
 	KUNIT_CASE(test_slab_buckets),
+	KUNIT_CASE(test_kmalloc_partition),
 #ifdef CONFIG_PERF_EVENTS
 	KUNIT_CASE_SLOW(test_kmalloc_nolock_and_friends_perf),
 #endif
diff --git a/mm/slab.h b/mm/slab.h
index 6a28951f6d5d..2a9d6a9a9948 100644
--- a/mm/slab.h
+++ b/mm/slab.h
@@ -391,6 +391,7 @@ extern struct list_head slab_caches;
 extern struct kmem_cache *kmem_cache;
 
 extern bool slab_buckets_enabled;
+extern bool kmalloc_partition_enabled;
 
 /* A table of kmalloc cache names and sizes */
 extern const struct kmalloc_info_struct {
diff --git a/mm/slab_common.c b/mm/slab_common.c
index 20b066d61703..d0695e4d2ab6 100644
--- a/mm/slab_common.c
+++ b/mm/slab_common.c
@@ -781,8 +781,9 @@ unsigned long random_kmalloc_seed __ro_after_init;
 EXPORT_SYMBOL(random_kmalloc_seed);
 #endif
 
-static bool kmalloc_partition_enabled __ro_after_init =
+bool kmalloc_partition_enabled __ro_after_init =
 	IS_ENABLED(CONFIG_KMALLOC_PARTITION_CACHES);
+EXPORT_SYMBOL_IF_KUNIT(kmalloc_partition_enabled);
 
 static int __init setup_kmalloc_partition(char *str)
 {
-- 
2.56.1