1 files changed, 6 insertions, 1 deletions

diff --git a/kernel/cpu.c b/kernel/cpu.c
index 8ce10043e4ac..ff071e022a85 100644
--- a/kernel/cpu.c
+++ b/kernel/cpu.c
@@ -14,6 +14,7 @@
 #include <linux/kthread.h>
 #include <linux/stop_machine.h>
 #include <linux/mutex.h>
+#include <asm/tboot.h>
 
 #ifdef CONFIG_SMP
 /* Serializes the updates to cpu_online_mask, cpu_present_mask */
@@ -376,7 +377,7 @@ static cpumask_var_t frozen_cpus;
 
 int disable_nonboot_cpus(void)
 {
-        int cpu, first_cpu, error;
+        int cpu, first_cpu, error, num_cpus = 0;
 
         error = stop_machine_create();
         if (error)
@@ -391,6 +392,7 @@ int disable_nonboot_cpus(void)
         for_each_online_cpu(cpu) {
                 if (cpu == first_cpu)
                         continue;
+                num_cpus++;
                 error = _cpu_down(cpu, 1);
                 if (!error) {
                         cpumask_set_cpu(cpu, frozen_cpus);
@@ -401,6 +403,9 @@ int disable_nonboot_cpus(void)
                         break;
                 }
         }
+        /* ensure all CPUs have gone into wait-for-SIPI */
+        error |= tboot_wait_for_aps(num_cpus);
+
         if (!error) {
                 BUG_ON(num_online_cpus() > 1);
                 /* Make sure the CPUs won't be enabled by someone else */