[PATCH v2] mm/memcontrol: skip non-hierarchical memcg-wide stats on the default hierarchy

[Joanne Koong <joannelkoong@gmail.com>]

memcg_vmstats keeps a non-hierarchical copy of every memcg-wide stat item
and event alongside the hierarchical one. The only readers however are
the legacy memory.stat and memory.numa_stat, and reparenting on offline,
which are only used when on the legacy v1 hierarchy.

On the default hierarchy, memcg_vmstats's non-hierarchical arrays are
written to on every rstat flush, despite their values never being read /
accessed.

Instead, skip non-hierarchical memcg-wide stats on the default hierarchy.
This makes flushes cheaper. mem_cgroup_stat_aggregate() can now skip the
read-modify-write of ac->local[i]. When on the default hierarchy,
nothing else accesses state_local or events_local, so those cachelines
were getting pulled in solely for the writes, and they are separate from
the ones the loop is already walking / accessing.

There is the unlikely case where a controller can be moved from the
default hierarchy to the legacy v1 hierarchy at runtime, which means the
non-hierarchical memcg-wide arrays will need to be populated with the
correct values. This can be handled through the bind callback, which
gets called by rebind_subsystems(). We can copy over the counters from
the hierarchical arrays since rebinding is only allowed when the root is
the only memcg, so the hierarchical arrays and the non-hierarchical
arrays values should be the same.

On an 80-cpu x86_64 machine with 500 cgroups each running a workload that
dirties anon, file, dirty/writeback, slab, kmem, mlock, and reclaim
counters, timing mem_cgroup_css_rstat_flush() in-kernel in TSC ticks per
flush showed roughly

                         before   after    delta
memcg-wide aggregation   1231     1180     -4.1%
overall flush function   2452     2397     -2.2%

These numbers are from taking the median of 70 samples, one per 20s
window on each kernel. The 95% intervals observed on the two deltas are
[-5.41%, -1.76%] and [-4.53%, -0.14%]. The values above include the
timing overhead itself, so only the delta is meaningful here.

Counting the items that actually changed, a median of 1.5 of the 77
memcg-wide items (57 state + 20 events) had a non-zero per-cpu delta at
each flush, which means the benchmarks above are with one or two fewer
cachelines pulled in per flush. The count is low because the benchmark
reads memory.stat in a loop to keep the flush rate up. For cases where
flushes are triggered only by the 2s periodic worker, more changes will
have accumulated between flushes, so more cachelines are skipped and the
per-flush saving should be larger.

Signed-off-by: Joanne Koong <joannelkoong@gmail.com>
---
v1: https://lore.kernel.org/linux-mm/20260901232834.22221-1-joannelkoong@gmail.com/
Changes since v1:
  * Change from gating on builds w/out CONFIG_MEMCG_V1 to gating on the
    default hierarchy so CONFIG_MEMCG_V1=y kernels that do not use v1 benefit
    too (Yosry)

 mm/memcontrol.c | 89 +++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 82 insertions(+), 7 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 7ce50bccf126..196e1791c10d 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -687,6 +687,29 @@ struct memcg_vmstats {
 	atomic_long_t		stats_updates;
 };
 
+/*
+ * The non-hierarchical memcg-wide counters are read back only by the legacy
+ * memory.stat and by reparenting on offline, both of which are v1-only.
+ * They are skipped if the controller sits on the default hierarchy
+ * (mem_cgroup_bind() populates them if it is later moved onto the legacy
+ * hierarchy).
+ */
+static long *memcg_state_local_array(struct mem_cgroup *memcg)
+{
+	if (cgroup_subsys_on_dfl(memory_cgrp_subsys))
+		return NULL;
+
+	return memcg->vmstats->state_local;
+}
+
+static unsigned long *memcg_events_local_array(struct mem_cgroup *memcg)
+{
+	if (cgroup_subsys_on_dfl(memory_cgrp_subsys))
+		return NULL;
+
+	return memcg->vmstats->events_local;
+}
+
 /*
  * memcg and lruvec stats flushing
  *
@@ -4468,7 +4491,10 @@ static void mem_cgroup_css_reset(struct cgroup_subsys_state *css)
 struct aggregate_control {
 	/* pointer to the aggregated (CPU and subtree aggregated) counters */
 	long *aggregate;
-	/* pointer to the non-hierarchichal (CPU aggregated) counters */
+	/*
+	 * pointer to the non-hierarchical (CPU aggregated) counters or NULL to
+	 * skip updating them (see memcg_state_local_array())
+	 */
 	long *local;
 	/* pointer to the pending child counters during tree propagation */
 	long *pending;
@@ -4507,7 +4533,7 @@ static void mem_cgroup_stat_aggregate(struct aggregate_control *ac)
 		}
 
 		/* Aggregate counts on this level and propagate upwards */
-		if (delta_cpu)
+		if (delta_cpu && ac->local)
 			ac->local[i] += delta_cpu;
 
 		if (delta) {
@@ -4521,6 +4547,7 @@ static void mem_cgroup_stat_aggregate(struct aggregate_control *ac)
 #ifdef CONFIG_MEMCG_NMI_SAFETY_REQUIRES_ATOMIC
 static void flush_nmi_stats(struct mem_cgroup *memcg, struct mem_cgroup *parent)
 {
+	long *state_local = memcg_state_local_array(memcg);
 	int nid;
 
 	if (atomic_read(&memcg->kmem_stat)) {
@@ -4528,7 +4555,8 @@ static void flush_nmi_stats(struct mem_cgroup *memcg, struct mem_cgroup *parent)
 		int index = memcg_stats_index(MEMCG_KMEM);
 
 		memcg->vmstats->state[index] += kmem;
-		memcg->vmstats->state_local[index] += kmem;
+		if (state_local)
+			state_local[index] += kmem;
 		if (parent)
 			parent->vmstats->state_pending[index] += kmem;
 	}
@@ -4550,7 +4578,8 @@ static void flush_nmi_stats(struct mem_cgroup *memcg, struct mem_cgroup *parent)
 			if (plstats)
 				plstats->state_pending[index] += slab;
 			memcg->vmstats->state[index] += slab;
-			memcg->vmstats->state_local[index] += slab;
+			if (state_local)
+				state_local[index] += slab;
 			if (parent)
 				parent->vmstats->state_pending[index] += slab;
 		}
@@ -4563,7 +4592,8 @@ static void flush_nmi_stats(struct mem_cgroup *memcg, struct mem_cgroup *parent)
 			if (plstats)
 				plstats->state_pending[index] += slab;
 			memcg->vmstats->state[index] += slab;
-			memcg->vmstats->state_local[index] += slab;
+			if (state_local)
+				state_local[index] += slab;
 			if (parent)
 				parent->vmstats->state_pending[index] += slab;
 		}
@@ -4588,7 +4618,7 @@ static void mem_cgroup_css_rstat_flush(struct cgroup_subsys_state *css, int cpu)
 
 	ac = (struct aggregate_control) {
 		.aggregate = memcg->vmstats->state,
-		.local = memcg->vmstats->state_local,
+		.local = memcg_state_local_array(memcg),
 		.pending = memcg->vmstats->state_pending,
 		.ppending = parent ? parent->vmstats->state_pending : NULL,
 		.cstat = statc->state,
@@ -4599,7 +4629,7 @@ static void mem_cgroup_css_rstat_flush(struct cgroup_subsys_state *css, int cpu)
 
 	ac = (struct aggregate_control) {
 		.aggregate = memcg->vmstats->events,
-		.local = memcg->vmstats->events_local,
+		.local = memcg_events_local_array(memcg),
 		.pending = memcg->vmstats->events_pending,
 		.ppending = parent ? parent->vmstats->events_pending : NULL,
 		.cstat = statc->events,
@@ -5183,6 +5213,50 @@ static struct cftype memory_files[] = {
 	{ }	/* terminate */
 };
 
+#ifdef CONFIG_MEMCG_V1
+/*
+ * Called after the controller has moved between hierarchies, with the on_dfl
+ * key already in its new state.
+ *
+ * If the controller is moving from the default hierarchy to the legacy
+ * hierarchy, vmstats's state_local and events_local arrays need to be populated
+ * since they get read by the legacy hierarchy. Those values can just be taken
+ * from the vmstats's state and events arrays since rebind_subsystems() only
+ * allows the move when the root is the only memcg.
+ */
+static void mem_cgroup_bind(struct cgroup_subsys_state *root_css)
+{
+	struct mem_cgroup *memcg = mem_cgroup_from_css(root_css);
+	int i;
+
+	if (cgroup_subsys_on_dfl(memory_cgrp_subsys))
+		return;
+
+	/*
+	 * The flush is necessary because there might be descendants that got
+	 * destroyed right before the rebind that may have left counts in the
+	 * pending arrays that haven't yet been folded into the state and events
+	 * arrays
+	 */
+	__mem_cgroup_flush_stats(memcg, true);
+
+	/*
+	 * Not serialized against a concurrent flush. The periodic flusher runs
+	 * on root_mem_cgroup without cgroup_mutex and a flush increments state
+	 * and state_local together. If the flush happens in between when we
+	 * read state and write to state_local, its state_local increment is
+	 * overwritten and the counter stays short, but the difference would be
+	 * a single CPU's accumulated charges since the last flush and the
+	 * counters are approximate values.
+	 */
+	for (i = 0; i < MEMCG_VMSTAT_SIZE; i++)
+		memcg->vmstats->state_local[i] = memcg->vmstats->state[i];
+
+	for (i = 0; i < NR_MEMCG_EVENTS; i++)
+		memcg->vmstats->events_local[i] = memcg->vmstats->events[i];
+}
+#endif /* CONFIG_MEMCG_V1 */
+
 struct cgroup_subsys memory_cgrp_subsys = {
 	.css_alloc = mem_cgroup_css_alloc,
 	.css_online = mem_cgroup_css_online,
@@ -5196,6 +5270,7 @@ struct cgroup_subsys memory_cgrp_subsys = {
 	.exit = mem_cgroup_exit,
 	.dfl_cftypes = memory_files,
 #ifdef CONFIG_MEMCG_V1
+	.bind = mem_cgroup_bind,
 	.legacy_cftypes = mem_cgroup_legacy_files,
 #endif
 	.early_init = 0,
-- 
2.52.0