[PATCH 0/5] Enhance Union-Find with KUnit tests and optimization improvements

[PATCH 0/5] Enhance Union-Find with KUnit tests and optimization improvements

[Kuan-Wei Chiu]

This patch series adds KUnit tests for the Union-Find implementation
and optimizes the path compression in the uf_find() function to achieve
a lower tree height and improved efficiency. Additionally, it modifies
uf_union() to return a boolean value indicating whether a merge
occurred, enhancing the process of calculating the number of groups in
the cgroup cpuset.

Regards,
Kuan-Wei

Kuan-Wei Chiu (5):
  lib/union_find: Add EXPORT_SYMBOL() for uf_find() and uf_union()
  lib/union_find: Change uf_union() return type to bool
  lib: Add KUnit tests for Union-Find implementation
  lib/union_find: Optimize uf_find() with enhanced path compression
  cgroup/cpuset: Optimize total domain count using updated uf_union()

 MAINTAINERS                |  1 +
 include/linux/union_find.h |  2 +-
 kernel/cgroup/cpuset.c     | 10 ++----
 lib/Kconfig.debug          | 12 +++++++
 lib/Makefile               |  1 +
 lib/union_find.c           | 20 ++++++++---
 lib/union_find_kunit.c     | 74 ++++++++++++++++++++++++++++++++++++++
 7 files changed, 108 insertions(+), 12 deletions(-)
 create mode 100644 lib/union_find_kunit.c

-- 
2.34.1

[PATCH 1/5] lib/union_find: Add EXPORT_SYMBOL() for uf_find() and uf_union()

[Kuan-Wei Chiu]

Add EXPORT_SYMBOL() for the uf_find() and uf_union() functions to allow
kernel modules, including the KUnit tests for the Union-Find data
structure, to use these functions. This enhances the usability of the
Union-Find implementation in a modular context, facilitating easier
testing and integration.

Signed-off-by: Kuan-Wei Chiu <visitorckw@gmail.com>
---
 lib/union_find.c | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/lib/union_find.c b/lib/union_find.c
index 413b0f8adf7a..c9fd30b8059c 100644
--- a/lib/union_find.c
+++ b/lib/union_find.c
@@ -1,4 +1,5 @@
 // SPDX-License-Identifier: GPL-2.0
+#include <linux/export.h>
 #include <linux/union_find.h>
 
 /**
@@ -21,6 +22,7 @@ struct uf_node *uf_find(struct uf_node *node)
 	}
 	return node;
 }
+EXPORT_SYMBOL(uf_find);
 
 /**
  * uf_union - Merge two sets, using union by rank
@@ -47,3 +49,4 @@ void uf_union(struct uf_node *node1, struct uf_node *node2)
 		root1->rank++;
 	}
 }
+EXPORT_SYMBOL(uf_union);
-- 
2.34.1

[PATCH 2/5] lib/union_find: Change uf_union() return type to bool

[Kuan-Wei Chiu]

Modify the uf_union() function to return a bool indicating whether a
merge occurred. If the two nodes belong to the same set, the function
returns false, indicating no merge took place. Otherwise, it completes
the merge and returns true. This change allows the caller to easily
determine the number of distinct groups by tracking successful merges,
enhancing the usability of the Union-Find implementation.

Signed-off-by: Kuan-Wei Chiu <visitorckw@gmail.com>
---
 include/linux/union_find.h | 2 +-
 lib/union_find.c           | 8 ++++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/include/linux/union_find.h b/include/linux/union_find.h
index cfd49263c138..45c1a6fc6574 100644
--- a/include/linux/union_find.h
+++ b/include/linux/union_find.h
@@ -36,6 +36,6 @@ static inline void uf_node_init(struct uf_node *node)
 struct uf_node *uf_find(struct uf_node *node);
 
 /* Merge two intersecting nodes */
-void uf_union(struct uf_node *node1, struct uf_node *node2);
+bool uf_union(struct uf_node *node1, struct uf_node *node2);
 
 #endif /* __LINUX_UNION_FIND_H */
diff --git a/lib/union_find.c b/lib/union_find.c
index c9fd30b8059c..a20678da0220 100644
--- a/lib/union_find.c
+++ b/lib/union_find.c
@@ -31,14 +31,16 @@ EXPORT_SYMBOL(uf_find);
  *
  * This function merges the sets containing node1 and node2, by comparing
  * the ranks to keep the tree balanced.
+ *
+ * Returns true if the sets were merged, false if they were already in the same set.
  */
-void uf_union(struct uf_node *node1, struct uf_node *node2)
+bool uf_union(struct uf_node *node1, struct uf_node *node2)
 {
 	struct uf_node *root1 = uf_find(node1);
 	struct uf_node *root2 = uf_find(node2);
 
 	if (root1 == root2)
-		return;
+		return false;
 
 	if (root1->rank < root2->rank) {
 		root1->parent = root2;
@@ -48,5 +50,7 @@ void uf_union(struct uf_node *node1, struct uf_node *node2)
 		root2->parent = root1;
 		root1->rank++;
 	}
+
+	return true;
 }
 EXPORT_SYMBOL(uf_union);
-- 
2.34.1

[PATCH 3/5] lib: Add KUnit tests for Union-Find implementation

[Kuan-Wei Chiu]

Introduce a KUnit test suite for the Union-Find data structure. The
tests verify the functionality and correctness of the union and find
operations, including edge cases such as handling duplicate unions.
The addition of KUnit tests enhances the robustness of the Union-Find
implementation by ensuring its correctness under various scenarios.

Signed-off-by: Kuan-Wei Chiu <visitorckw@gmail.com>
---
Regarding the changes to the MAINTAINERS file, I'm also happy to help
maintain/review patches related to union find. If I am qualified
enough, may I send another patch to add myself later? :)

 MAINTAINERS            |  1 +
 lib/Kconfig.debug      | 12 +++++++
 lib/Makefile           |  1 +
 lib/union_find_kunit.c | 74 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 88 insertions(+)
 create mode 100644 lib/union_find_kunit.c

diff --git a/MAINTAINERS b/MAINTAINERS
index 5153c995d429..3b10ac1cdf63 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -23799,6 +23799,7 @@ F:	Documentation/core-api/union_find.rst
 F:	Documentation/translations/zh_CN/core-api/union_find.rst
 F:	include/linux/union_find.h
 F:	lib/union_find.c
+F:	lib/union_find_kunit.c
 
 UNIVERSAL FLASH STORAGE HOST CONTROLLER DRIVER
 R:	Alim Akhtar <alim.akhtar@samsung.com>
diff --git a/lib/Kconfig.debug b/lib/Kconfig.debug
index 7315f643817a..376c86d34253 100644
--- a/lib/Kconfig.debug
+++ b/lib/Kconfig.debug
@@ -2841,6 +2841,18 @@ config SIPHASH_KUNIT_TEST
 	  This is intended to help people writing architecture-specific
 	  optimized versions.  If unsure, say N.
 
+config UNION_FIND_KUNIT_TEST
+	tristate "KUnit Test for Union find"
+	depends on KUNIT
+	default KUNIT_ALL_TESTS
+	help
+	  This option enables the KUnit tests for the Union-Find data structure.
+	  These tests verify the functionality and correctness of the Union-Find
+	  implementation, including union and find operations, as well as
+	  edge cases such as handling of duplicate unions.
+
+	  If unsure, say N
+
 config USERCOPY_KUNIT_TEST
 	tristate "KUnit Test for user/kernel boundary protections"
 	depends on KUNIT
diff --git a/lib/Makefile b/lib/Makefile
index 773adf88af41..03da92faf9b8 100644
--- a/lib/Makefile
+++ b/lib/Makefile
@@ -388,6 +388,7 @@ CFLAGS_fortify_kunit.o += $(call cc-disable-warning, stringop-truncation)
 CFLAGS_fortify_kunit.o += $(DISABLE_STRUCTLEAK_PLUGIN)
 obj-$(CONFIG_FORTIFY_KUNIT_TEST) += fortify_kunit.o
 obj-$(CONFIG_SIPHASH_KUNIT_TEST) += siphash_kunit.o
+obj-$(CONFIG_UNION_FIND_KUNIT_TEST) += union_find_kunit.o
 obj-$(CONFIG_USERCOPY_KUNIT_TEST) += usercopy_kunit.o
 
 obj-$(CONFIG_GENERIC_LIB_DEVMEM_IS_ALLOWED) += devmem_is_allowed.o
diff --git a/lib/union_find_kunit.c b/lib/union_find_kunit.c
new file mode 100644
index 000000000000..9bdf9e0e455e
--- /dev/null
+++ b/lib/union_find_kunit.c
@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: GPL-2.0-only
+
+#include <kunit/test.h>
+#include <linux/module.h>
+#include <linux/union_find.h>
+
+static void test_union_and_find(struct kunit *test)
+{
+	struct uf_node node1, node2, node3;
+	struct uf_node *root1, *root2, *root3;
+	bool merged;
+
+	/* Initialize the nodes */
+	uf_node_init(&node1);
+	uf_node_init(&node2);
+	uf_node_init(&node3);
+
+	/* Check the initial parent and rank */
+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node1), &node1);
+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node2), &node2);
+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node3), &node3);
+	KUNIT_ASSERT_EQ(test, node1.rank, 0);
+	KUNIT_ASSERT_EQ(test, node2.rank, 0);
+	KUNIT_ASSERT_EQ(test, node3.rank, 0);
+
+	/* Union node1 and node2 */
+	merged = uf_union(&node1, &node2);
+	KUNIT_ASSERT_TRUE(test, merged);
+
+	/* Assert that one of the nodes is now the parent of the other */
+	root1 = uf_find(&node1);
+	root2 = uf_find(&node2);
+	KUNIT_ASSERT_PTR_EQ(test, root1, root2);
+
+	/* Check rank after the first union */
+	if (root1 == &node1) {
+		KUNIT_ASSERT_EQ(test, node1.rank, 1);
+		KUNIT_ASSERT_EQ(test, node2.rank, 0);
+	} else {
+		KUNIT_ASSERT_EQ(test, node1.rank, 0);
+		KUNIT_ASSERT_EQ(test, node2.rank, 1);
+	}
+
+	/* Attempt to union node1 and node2 again and check for false return */
+	merged = uf_union(&node1, &node2);
+	KUNIT_ASSERT_FALSE(test, merged);
+
+	/* Union node3 with the result of the previous union (node1 and node2) */
+	uf_union(&node1, &node3);
+
+	/* Assert that all nodes have the same root */
+	root3 = uf_find(&node3);
+	KUNIT_ASSERT_PTR_EQ(test, root1, root3);
+
+	/* Check rank after the second union */
+	KUNIT_ASSERT_EQ(test, root1->rank, 1);
+	KUNIT_ASSERT_EQ(test, node3.rank, 0);
+}
+
+static struct kunit_case union_find_test_cases[] = {
+	KUNIT_CASE(test_union_and_find),
+	{}
+};
+
+static struct kunit_suite union_find_test_suite = {
+	.name = "union_find_test_suite",
+	.test_cases = union_find_test_cases,
+};
+
+kunit_test_suites(&union_find_test_suite);
+
+MODULE_AUTHOR("Kuan-Wei Chiu <visitorckw@gmail.com>");
+MODULE_DESCRIPTION("Union-find KUnit test suite");
+MODULE_LICENSE("GPL");
-- 
2.34.1

[PATCH 4/5] lib/union_find: Optimize uf_find() with enhanced path compression

[Kuan-Wei Chiu]

Optimize the uf_find() function to enhance its efficiency by
implementing a more effective path compression strategy. The original
implementation only updated the parent pointer of the current node to
its grandparent, resulting in a relatively shallow tree.

In the updated version, once the root of the node is identified, all
nodes along the search path are updated to directly point to the root.
This change minimizes the height of the tree and improves the
efficiency for subsequent find operations, providing better performance
for the Union-Find data structure.

Signed-off-by: Kuan-Wei Chiu <visitorckw@gmail.com>
---
Note: Tested with the KUnit tests introduced in the previous patch.

 lib/union_find.c | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/lib/union_find.c b/lib/union_find.c
index a20678da0220..7c553fa622c8 100644
--- a/lib/union_find.c
+++ b/lib/union_find.c
@@ -13,14 +13,19 @@
  */
 struct uf_node *uf_find(struct uf_node *node)
 {
+	struct uf_node *root = node;
 	struct uf_node *parent;
 
+	while (root->parent != root)
+		root = root->parent;
+
 	while (node->parent != node) {
 		parent = node->parent;
-		node->parent = parent->parent;
+		node->parent = root;
 		node = parent;
 	}
-	return node;
+
+	return root;
 }
 EXPORT_SYMBOL(uf_find);
 
-- 
2.34.1

[PATCH 5/5] cgroup/cpuset: Optimize total domain count using updated uf_union()

[Kuan-Wei Chiu]

Improve the efficiency of calculating the total number of scheduling
domains by using the updated uf_union function, which now returns a
boolean to indicate if a merge occurred. Previously, an additional loop
was needed to count root nodes for distinct groups. With this change,
each successful merge reduces the domain count (ndoms) directly,
eliminating the need for the final loop and enhancing performance.

Signed-off-by: Kuan-Wei Chiu <visitorckw@gmail.com>
---
Note: Tested with test_cpuset_prs.sh

 kernel/cgroup/cpuset.c | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/kernel/cgroup/cpuset.c b/kernel/cgroup/cpuset.c
index a4dd285cdf39..5e9301550d43 100644
--- a/kernel/cgroup/cpuset.c
+++ b/kernel/cgroup/cpuset.c
@@ -817,6 +817,8 @@ static int generate_sched_domains(cpumask_var_t **domains,
 	if (root_load_balance && (csn == 1))
 		goto single_root_domain;
 
+	ndoms = csn;
+
 	for (i = 0; i < csn; i++)
 		uf_node_init(&csa[i]->node);
 
@@ -829,17 +831,11 @@ static int generate_sched_domains(cpumask_var_t **domains,
 				 * partition root cpusets.
 				 */
 				WARN_ON_ONCE(cgrpv2);
-				uf_union(&csa[i]->node, &csa[j]->node);
+				ndoms -= uf_union(&csa[i]->node, &csa[j]->node);
 			}
 		}
 	}
 
-	/* Count the total number of domains */
-	for (i = 0; i < csn; i++) {
-		if (uf_find(&csa[i]->node) == &csa[i]->node)
-			ndoms++;
-	}
-
 	/*
 	 * Now we know how many domains to create.
 	 * Convert <csn, csa> to <ndoms, doms> and populate cpu masks.
-- 
2.34.1

Re:[PATCH 3/5] lib: Add KUnit tests for Union-Find implementation

[Xavier]

At 2024-10-06 05:49:36, "Kuan-Wei Chiu" <visitorckw@gmail.com> wrote:
>Introduce a KUnit test suite for the Union-Find data structure. The
>tests verify the functionality and correctness of the union and find
>operations, including edge cases such as handling duplicate unions.
>The addition of KUnit tests enhances the robustness of the Union-Find
>implementation by ensuring its correctness under various scenarios.
>
>Signed-off-by: Kuan-Wei Chiu <visitorckw@gmail.com>
>---
>Regarding the changes to the MAINTAINERS file, I'm also happy to help
>maintain/review patches related to union find. If I am qualified
>enough, may I send another patch to add myself later? :)

Of course, if you are interested.

>
> MAINTAINERS            |  1 +
> lib/Kconfig.debug      | 12 +++++++
> lib/Makefile           |  1 +
> lib/union_find_kunit.c | 74 ++++++++++++++++++++++++++++++++++++++++++
> 4 files changed, 88 insertions(+)
> create mode 100644 lib/union_find_kunit.c
>
>diff --git a/MAINTAINERS b/MAINTAINERS
>index 5153c995d429..3b10ac1cdf63 100644
>--- a/MAINTAINERS
>+++ b/MAINTAINERS
>@@ -23799,6 +23799,7 @@ F:	Documentation/core-api/union_find.rst
> F:	Documentation/translations/zh_CN/core-api/union_find.rst
> F:	include/linux/union_find.h
> F:	lib/union_find.c
>+F:	lib/union_find_kunit.c
> 
> UNIVERSAL FLASH STORAGE HOST CONTROLLER DRIVER
> R:	Alim Akhtar <alim.akhtar@samsung.com>
>diff --git a/lib/Kconfig.debug b/lib/Kconfig.debug
>index 7315f643817a..376c86d34253 100644
>--- a/lib/Kconfig.debug
>+++ b/lib/Kconfig.debug
>@@ -2841,6 +2841,18 @@ config SIPHASH_KUNIT_TEST
> 	  This is intended to help people writing architecture-specific
> 	  optimized versions.  If unsure, say N.
> 
>+config UNION_FIND_KUNIT_TEST
>+	tristate "KUnit Test for Union find"
>+	depends on KUNIT
>+	default KUNIT_ALL_TESTS
>+	help
>+	  This option enables the KUnit tests for the Union-Find data structure.
>+	  These tests verify the functionality and correctness of the Union-Find
>+	  implementation, including union and find operations, as well as
>+	  edge cases such as handling of duplicate unions.
>+
>+	  If unsure, say N
>+
> config USERCOPY_KUNIT_TEST
> 	tristate "KUnit Test for user/kernel boundary protections"
> 	depends on KUNIT
>diff --git a/lib/Makefile b/lib/Makefile
>index 773adf88af41..03da92faf9b8 100644
>--- a/lib/Makefile
>+++ b/lib/Makefile
>@@ -388,6 +388,7 @@ CFLAGS_fortify_kunit.o += $(call cc-disable-warning, stringop-truncation)
> CFLAGS_fortify_kunit.o += $(DISABLE_STRUCTLEAK_PLUGIN)
> obj-$(CONFIG_FORTIFY_KUNIT_TEST) += fortify_kunit.o
> obj-$(CONFIG_SIPHASH_KUNIT_TEST) += siphash_kunit.o
>+obj-$(CONFIG_UNION_FIND_KUNIT_TEST) += union_find_kunit.o
> obj-$(CONFIG_USERCOPY_KUNIT_TEST) += usercopy_kunit.o
> 
> obj-$(CONFIG_GENERIC_LIB_DEVMEM_IS_ALLOWED) += devmem_is_allowed.o
>diff --git a/lib/union_find_kunit.c b/lib/union_find_kunit.c
>new file mode 100644
>index 000000000000..9bdf9e0e455e
>--- /dev/null
>+++ b/lib/union_find_kunit.c
>@@ -0,0 +1,74 @@
>+// SPDX-License-Identifier: GPL-2.0-only
>+
>+#include <kunit/test.h>
>+#include <linux/module.h>
>+#include <linux/union_find.h>
>+
>+static void test_union_and_find(struct kunit *test)
>+{
>+	struct uf_node node1, node2, node3;
>+	struct uf_node *root1, *root2, *root3;
>+	bool merged;
>+
>+	/* Initialize the nodes */
>+	uf_node_init(&node1);
>+	uf_node_init(&node2);
>+	uf_node_init(&node3);
>+
>+	/* Check the initial parent and rank */
>+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node1), &node1);
>+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node2), &node2);
>+	KUNIT_ASSERT_PTR_EQ(test, uf_find(&node3), &node3);
>+	KUNIT_ASSERT_EQ(test, node1.rank, 0);
>+	KUNIT_ASSERT_EQ(test, node2.rank, 0);
>+	KUNIT_ASSERT_EQ(test, node3.rank, 0);
>+
>+	/* Union node1 and node2 */
>+	merged = uf_union(&node1, &node2);
>+	KUNIT_ASSERT_TRUE(test, merged);
>+
>+	/* Assert that one of the nodes is now the parent of the other */
>+	root1 = uf_find(&node1);
>+	root2 = uf_find(&node2);
>+	KUNIT_ASSERT_PTR_EQ(test, root1, root2);
>+
>+	/* Check rank after the first union */
>+	if (root1 == &node1) {
>+		KUNIT_ASSERT_EQ(test, node1.rank, 1);
>+		KUNIT_ASSERT_EQ(test, node2.rank, 0);
>+	} else {
>+		KUNIT_ASSERT_EQ(test, node1.rank, 0);
>+		KUNIT_ASSERT_EQ(test, node2.rank, 1);
>+	}
>+
>+	/* Attempt to union node1 and node2 again and check for false return */
>+	merged = uf_union(&node1, &node2);
>+	KUNIT_ASSERT_FALSE(test, merged);
>+
>+	/* Union node3 with the result of the previous union (node1 and node2) */
>+	uf_union(&node1, &node3);
>+
>+	/* Assert that all nodes have the same root */
>+	root3 = uf_find(&node3);
>+	KUNIT_ASSERT_PTR_EQ(test, root1, root3);
>+
>+	/* Check rank after the second union */
>+	KUNIT_ASSERT_EQ(test, root1->rank, 1);
>+	KUNIT_ASSERT_EQ(test, node3.rank, 0);
>+}
>+
>+static struct kunit_case union_find_test_cases[] = {
>+	KUNIT_CASE(test_union_and_find),
>+	{}
>+};
>+
>+static struct kunit_suite union_find_test_suite = {
>+	.name = "union_find_test_suite",
>+	.test_cases = union_find_test_cases,
>+};
>+
>+kunit_test_suites(&union_find_test_suite);
>+
>+MODULE_AUTHOR("Kuan-Wei Chiu <visitorckw@gmail.com>");
>+MODULE_DESCRIPTION("Union-find KUnit test suite");
>+MODULE_LICENSE("GPL");
>-- 
>2.34.1

Re:[PATCH 4/5] lib/union_find: Optimize uf_find() with enhanced path compression

[Xavier]

At 2024-10-06 05:49:37, "Kuan-Wei Chiu" <visitorckw@gmail.com> wrote:
>Optimize the uf_find() function to enhance its efficiency by
>implementing a more effective path compression strategy. The original
>implementation only updated the parent pointer of the current node to
>its grandparent, resulting in a relatively shallow tree.
>
>In the updated version, once the root of the node is identified, all
>nodes along the search path are updated to directly point to the root.
>This change minimizes the height of the tree and improves the
>efficiency for subsequent find operations, providing better performance
>for the Union-Find data structure.
>
>Signed-off-by: Kuan-Wei Chiu <visitorckw@gmail.com>
>---
>Note: Tested with the KUnit tests introduced in the previous patch.
>
> lib/union_find.c | 9 +++++++--
> 1 file changed, 7 insertions(+), 2 deletions(-)
>
>diff --git a/lib/union_find.c b/lib/union_find.c
>index a20678da0220..7c553fa622c8 100644
>--- a/lib/union_find.c
>+++ b/lib/union_find.c
>@@ -13,14 +13,19 @@
>  */
> struct uf_node *uf_find(struct uf_node *node)
> {
>+	struct uf_node *root = node;
> 	struct uf_node *parent;
> 
>+	while (root->parent != root)
>+		root = root->parent;
>+
> 	while (node->parent != node) {

Using “root” for this judgment might be better, as it could
reduce unnecessary entering.
        while (node->parent != root) {

> 		parent = node->parent;
>-		node->parent = parent->parent;
>+		node->parent = root;
> 		node = parent;
> 	}
>-	return node;
>+
>+	return root;
> }
> EXPORT_SYMBOL(uf_find);
> 
>-- 
>2.34.1