Arch specific code can update the mask based on the steal time. For
debugging it is desired to overwrite the arch logic. Do that with this
helper script.

Signed-off-by: Shrikanth Hegde <sshegde@linux.ibm.com>
---
 arch/powerpc/platforms/pseries/lpar.c | 3 +++
 drivers/base/cpu.c                    | 2 ++
 include/linux/sched.h                 | 4 ++++
 kernel/sched/core.c                   | 1 +
 4 files changed, 10 insertions(+)

diff --git a/arch/powerpc/platforms/pseries/lpar.c b/arch/powerpc/platforms/pseries/lpar.c
index 090e5c48243b..04bc75e22e7b 100644
--- a/arch/powerpc/platforms/pseries/lpar.c
+++ b/arch/powerpc/platforms/pseries/lpar.c
@@ -681,6 +681,9 @@ static void process_steal(int cpu)
 	unsigned long steal = 0;
 	unsigned int i;
 
+	if (static_branch_unlikely(&disable_arch_paravirt_handling))
+		return;
+
 	if (!should_cpu_process_steal(cpu))
 		return;
 
diff --git a/drivers/base/cpu.c b/drivers/base/cpu.c
index 043e4f4ce1a9..fbaddbfe0b01 100644
--- a/drivers/base/cpu.c
+++ b/drivers/base/cpu.c
@@ -402,7 +402,9 @@ static ssize_t store_paravirt_cpus(struct device *dev,
 	/* No more paravirt cpus */
 	if (cpumask_empty(temp_mask)) {
 		cpumask_copy((struct cpumask *)&__cpu_paravirt_mask, temp_mask);
+		static_branch_disable(&disable_arch_paravirt_handling);
 	} else {
+		static_branch_enable(&disable_arch_paravirt_handling);
 		cpumask_copy((struct cpumask *)&__cpu_paravirt_mask, temp_mask);
 
 		/* Enable tick on nohz_full cpu */
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 3628edd1468b..1afa5dd5b0ae 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -2427,4 +2427,8 @@ extern void migrate_enable(void);
 
 DEFINE_LOCK_GUARD_0(migrate, migrate_disable(), migrate_enable())
 
+#ifdef CONFIG_PARAVIRT
+DECLARE_STATIC_KEY_FALSE(disable_arch_paravirt_handling);
+#endif
+
 #endif
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 65c247c24191..b65a9898c694 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -10873,6 +10873,7 @@ void sched_change_end(struct sched_change_ctx *ctx)
 #ifdef CONFIG_PARAVIRT
 struct cpumask __cpu_paravirt_mask __read_mostly;
 EXPORT_SYMBOL(__cpu_paravirt_mask);
+DEFINE_STATIC_KEY_FALSE(disable_arch_paravirt_handling);
 
 static DEFINE_PER_CPU(struct cpu_stop_work, pv_push_task_work);
 
-- 
2.47.3