diff --git a/kernel/sched/walt/walt.c b/kernel/sched/walt/walt.c
index b479186d4f42..3b2aed8a82d4 100644
--- a/kernel/sched/walt/walt.c
+++ b/kernel/sched/walt/walt.c
@@ -4208,6 +4208,23 @@ static void dump_throttled_rt_tasks(void *unused, int cpu, u64 clock,
 	BUG_ON(sysctl_sched_bug_on_rt_throttle);
 }
 
+static void android_rvh_set_cpus_allowed_ptr_locked(void *unused,
+						    const struct cpumask *cpu_valid_mask,
+						    const struct cpumask *new_mask,
+						    unsigned int *dest_cpu)
+{
+	cpumask_t allowed_cpus;
+
+	if (unlikely(walt_disabled))
+		return;
+
+	if (cpu_halted(*dest_cpu)) {
+		/* remove halted cpus from the valid mask, and store locally */
+		cpumask_andnot(&allowed_cpus, cpu_valid_mask, cpu_halt_mask);
+		*dest_cpu = cpumask_any_and_distribute(&allowed_cpus, new_mask);
+	}
+}
+
 static void register_walt_hooks(void)
 {
 	register_trace_android_rvh_wake_up_new_task(android_rvh_wake_up_new_task, NULL);
@@ -4236,6 +4253,8 @@ static void register_walt_hooks(void)
 	register_trace_android_rvh_build_perf_domains(android_rvh_build_perf_domains, NULL);
 	register_trace_cpu_frequency_limits(walt_cpu_frequency_limits, NULL);
 	register_trace_android_vh_dump_throttled_rt_tasks(dump_throttled_rt_tasks, NULL);
+	register_trace_android_rvh_set_cpus_allowed_ptr_locked(
+						android_rvh_set_cpus_allowed_ptr_locked, NULL);
 }
 
 atomic64_t walt_irq_work_lastq_ws;