@@ -220,6 +220,49 @@ static bool test_state(unsigned int *tasks, enum psi_states state)
 	}
 }
 
+static u32 read_update_delta(struct psi_group_cpu *groupc,
+			     enum psi_states state, int cpu)
+{
+	u32 time, delta;
+
+	time = READ_ONCE(groupc->times[state]);
+	/*
+	 * In addition to already concluded states, we also
+	 * incorporate currently active states on the CPU, since
+	 * states may last for many sampling periods.
+	 *
+	 * This way we keep our delta sampling buckets small (u32) and
+	 * our reported pressure close to what's actually happening.
+	 */
+	if (test_state(groupc->tasks, state)) {
+		/*
+		 * We can race with a state change and need to make
+		 * sure the state_start update is ordered against the
+		 * updates to the live state and the time buckets
+		 * (groupc->times).
+		 *
+		 * 1. If we observe task state that needs to be
+		 * recorded, make sure we see state_start from when
+		 * that state went into effect or we'll count time
+		 * from the previous state.
+		 *
+		 * 2. If the time delta has already been added to the
+		 * bucket, make sure we don't see it in state_start or
+		 * we'll count it twice.
+		 *
+		 * If the time delta is out of state_start but not in
+		 * the time bucket yet, we'll miss it entirely and
+		 * handle it in the next period.
+		 */
+		smp_rmb();
+		time += cpu_clock(cpu) - groupc->state_start;
+	}
+	delta = time - groupc->times_prev[state];
+	groupc->times_prev[state] = time;
+
+	return delta;
+}
+
 static bool psi_update_stats(struct psi_group *group)
 {
 	u64 deltas[NR_PSI_STATES - 1] = { 0, };

@@ -244,60 +287,17 @@ static bool psi_update_stats(struct psi_group *group)
 	 */
 	for_each_online_cpu(cpu) {
 		struct psi_group_cpu *groupc = per_cpu_ptr(group->pcpu, cpu);
-		u32 uninitialized_var(nonidle);
-
-		BUILD_BUG_ON(PSI_NONIDLE != NR_PSI_STATES - 1);
-
-		for (s = PSI_NONIDLE; s >= 0; s--) {
-			u32 time, delta;
-
-			time = READ_ONCE(groupc->times[s]);
-			/*
-			 * In addition to already concluded states, we
-			 * also incorporate currently active states on
-			 * the CPU, since states may last for many
-			 * sampling periods.
-			 *
-			 * This way we keep our delta sampling buckets
-			 * small (u32) and our reported pressure close
-			 * to what's actually happening.
-			 */
-			if (test_state(groupc->tasks, cpu, s)) {
-				/*
-				 * We can race with a state change and
-				 * need to make sure the state_start
-				 * update is ordered against the
-				 * updates to the live state and the
-				 * time buckets (groupc->times).
-				 *
-				 * 1. If we observe task state that
-				 * needs to be recorded, make sure we
-				 * see state_start from when that
-				 * state went into effect or we'll
-				 * count time from the previous state.
-				 *
-				 * 2. If the time delta has already
-				 * been added to the bucket, make sure
-				 * we don't see it in state_start or
-				 * we'll count it twice.
-				 *
-				 * If the time delta is out of
-				 * state_start but not in the time
-				 * bucket yet, we'll miss it entirely
-				 * and handle it in the next period.
-				 */
-				smp_rmb();
-				time += cpu_clock(cpu) - groupc->state_start;
-			}
-			delta = time - groupc->times_prev[s];
-			groupc->times_prev[s] = time;
-
-			if (s == PSI_NONIDLE) {
-				nonidle = nsecs_to_jiffies(delta);
-				nonidle_total += nonidle;
-			} else {
-				deltas[s] += (u64)delta * nonidle;
-			}
+		u32 nonidle;
+
+		nonidle = read_update_delta(groupc, PSI_NONIDLE, cpu);
+		nonidle = nsecs_to_jiffies(nonidle);
+		nonidle_total += nonidle;
+
+		for (s = 0; s < PSI_NONIDLE; s++) {
+			u32 delta;
+
+			delta = read_update_delta(groupc, s, cpu);
+			deltas[s] += (u64)delta * nonidle;
 		}
 	}