[v11,6/7] microcode: rendezvous CPUs in NMI handler and load ucode

Commit Message

Chao Gao Sept. 26, 2019, 1:53 p.m. UTC

When one core is loading ucode, handling NMI on sibling threads or
on other cores in the system might be problematic. By rendezvousing
all CPUs in NMI handler, it prevents NMI acceptance during ucode
loading.

Basically, some work previously done in stop_machine context is
moved to NMI handler. Primary threads call in and load ucode in
NMI handler. Secondary threads wait for the completion of ucode
loading on all CPU cores. An option is introduced to disable this
behavior.

Control thread doesn't rendezvous in NMI handler by calling self_nmi()
(in case of unknown_nmi_error() being triggered). The side effect is
control thread might be handling an NMI and interacting with the old
ucode not in a controlled way while other threads are loading ucode.
Update ucode on the control thread first to mitigate this issue.

Signed-off-by: Sergey Dyasli <[e-mail]>
Signed-off-by: Chao Gao <[e-mail]>
---
Changes in v11:
 - Extend existing 'nmi' option rather than use a new one.
 - use per-cpu variable to store error code of xxx_nmi_work()
 - rename secondary_thread_work to secondary_nmi_work.
 - intialize nmi_patch to ZERO_BLOCK_PTR and make it static.
 - constify nmi_cpu
 - explain why control thread loads ucode first in patch description

Changes in v10:
 - rewrite based on Sergey's idea and patch
 - add Sergey's SOB.
 - add an option to disable ucode loading in NMI handler
 - don't send IPI NMI to the control thread to avoid unknown_nmi_error()
 in do_nmi().
 - add an assertion to make sure the cpu chosen to handle platform NMI
 won't send self NMI. Otherwise, there is a risk that we encounter
 unknown_nmi_error() and system crashes.

Changes in v9:
 - control threads send NMI to all other threads. Slave threads will
 stay in the NMI handling to prevent NMI acceptance during ucode
 loading. Note that self-nmi is invalid according to SDM.
 - s/rep_nop/cpu_relax
 - remove debug message in microcode_nmi_callback(). Printing debug
 message would take long times and control thread may timeout.
 - rebase and fix conflicts

Changes in v8:
 - new
---
 docs/misc/xen-command-line.pandoc |   6 +-
 xen/arch/x86/microcode.c          | 156 ++++++++++++++++++++++++++++++--------
 xen/arch/x86/traps.c              |   6 +-
 xen/include/asm-x86/nmi.h         |   3 +
 4 files changed, 138 insertions(+), 33 deletions(-)

Comments

Jan Beulich Sept. 27, 2019, 10:19 a.m. UTC | #1

On 26.09.2019 15:53, Chao Gao wrote:
> @@ -105,23 +110,42 @@ void __init microcode_set_module(unsigned int idx)
>  }
>  
>  /*
> - * The format is '[<integer>|scan]'. Both options are optional.
> + * The format is '[<integer>|scan, nmi=<bool>]'. Both options are optional.
>   * If the EFI has forced which of the multiboot payloads is to be used,
> - * no parsing will be attempted.
> + * only nmi=<bool> is parsed.
>   */
>  static int __init parse_ucode(const char *s)
>  {
> -    const char *q = NULL;
> +    const char *ss;
> +    int val, rc = 0;
>  
> -    if ( ucode_mod_forced ) /* Forced by EFI */
> -       return 0;
> +    do {
> +        ss = strchr(s, ',');
> +        if ( !ss )
> +            ss = strchr(s, '\0');
>  
> -    if ( !strncmp(s, "scan", 4) )
> -        ucode_scan = 1;
> -    else
> -        ucode_mod_idx = simple_strtol(s, &q, 0);
> +        if ( (val = parse_boolean("nmi", s, ss)) >= 0 )
> +            ucode_in_nmi = val;
> +        else if ( !ucode_mod_forced ) /* Not forced by EFI */
> +        {
> +            const char *q = NULL;
> +
> +            if ( !strncmp(s, "scan", 4) )
> +            {
> +                ucode_scan = true;

I guess it would have resulted in more consistent code if you had
used parse_boolean() here, too.

> @@ -222,6 +246,8 @@ const struct microcode_ops *microcode_ops;
>  static DEFINE_SPINLOCK(microcode_mutex);
>  
>  DEFINE_PER_CPU(struct cpu_signature, cpu_sig);
> +/* Store error code of the work done in NMI handler */
> +DEFINE_PER_CPU(int, loading_err);

static

> @@ -356,42 +383,88 @@ static void set_state(unsigned int state)
>      smp_wmb();
>  }
>  
> -static int secondary_thread_fn(void)
> +static int secondary_nmi_work(void)
>  {
> -    unsigned int primary = cpumask_first(this_cpu(cpu_sibling_mask));
> +    cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
>  
> -    if ( !wait_for_state(LOADING_CALLIN) )
> -        return -EBUSY;
> +    return wait_for_state(LOADING_EXIT) ? 0 : -EBUSY;
> +}
> +
> +static int primary_thread_work(const struct microcode_patch *patch)
> +{
> +    int ret;
>  
>      cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
>  
> -    if ( !wait_for_state(LOADING_EXIT) )
> +    if ( !wait_for_state(LOADING_ENTER) )
>          return -EBUSY;
>  
> -    /* Copy update revision from the primary thread. */
> -    this_cpu(cpu_sig).rev = per_cpu(cpu_sig, primary).rev;
> +    ret = microcode_ops->apply_microcode(patch);
> +    if ( !ret )
> +        atomic_inc(&cpu_updated);
> +    atomic_inc(&cpu_out);
>  
> -    return 0;
> +    return ret;
>  }
>  
> -static int primary_thread_fn(const struct microcode_patch *patch)
> +static int primary_nmi_work(const struct microcode_patch *patch)
> +{
> +    return primary_thread_work(patch);
> +}

Why this wrapper? The function signatures are identical. I guess
you want to emphasize the environment the function is to be used
in, so perhaps fine despite the redundancy. At least there's no
address taken of this function, so the compiler can eliminate it.

> +static int secondary_thread_fn(void)
> +{
>      if ( !wait_for_state(LOADING_CALLIN) )
>          return -EBUSY;
>  
> -    cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
> +    self_nmi();
>  
> -    if ( !wait_for_state(LOADING_ENTER) )
> +    /* Copy update revision from the primary thread. */
> +    this_cpu(cpu_sig).rev =
> +        per_cpu(cpu_sig, cpumask_first(this_cpu(cpu_sibling_mask))).rev;

_alternative_instructions() takes specific care to avoid relying on
the NMI potentially not arriving synchronously (in which case you'd
potentially copy a not-yet-updated CPU signature above). I think the
same care wants applying here, which I guess would be another

    wait_for_state(LOADING_EXIT);

> +    return this_cpu(loading_err);
> +}
> +
> +static int primary_thread_fn(const struct microcode_patch *patch)
> +{
> +    if ( !wait_for_state(LOADING_CALLIN) )
>          return -EBUSY;
>  
> -    ret = microcode_ops->apply_microcode(patch);
> -    if ( !ret )
> -        atomic_inc(&cpu_updated);
> -    atomic_inc(&cpu_out);
> +    if ( ucode_in_nmi )
> +    {
> +        self_nmi();
> +        return this_cpu(loading_err);

Same here than, to protect against returning a not-yet-updated error
indicator.

> @@ -420,14 +498,23 @@ static int control_thread_fn(const struct microcode_patch *patch)
>          return ret;
>      }
>  
> -    /* Let primary threads load the given ucode update */
> -    set_state(LOADING_ENTER);
> -
> +    /* Control thread loads ucode first while others are in NMI handler. */
>      ret = microcode_ops->apply_microcode(patch);
>      if ( !ret )
>          atomic_inc(&cpu_updated);
>      atomic_inc(&cpu_out);
>  
> +    if ( ret == -EIO )
> +    {
> +        printk(XENLOG_ERR
> +               "Late loading aborted: CPU%u failed to update ucode\n", cpu);
> +        set_state(LOADING_EXIT);
> +        return ret;
> +    }
> +
> +    /* Let primary threads load the given ucode update */
> +    set_state(LOADING_ENTER);

While the description goes to some lengths to explain this ordering of
updates, I still don't really see the point: How is it better for the
control CPU to have updated its ucode early and then hit an NMI before
the other CPUs have even started updating, than the other way around
in the opposite case?

> @@ -456,6 +543,8 @@ static int control_thread_fn(const struct microcode_patch *patch)
>      /* Mark loading is done to unblock other threads */
>      set_state(LOADING_EXIT);
>  
> +    set_nmi_callback(saved_nmi_callback);
> +    nmi_patch = ZERO_BLOCK_PTR;

Another smp_wmb() between them, just to be on the safe side?

> @@ -515,6 +604,13 @@ int microcode_update(XEN_GUEST_HANDLE_PARAM(const_void) buf, unsigned long len)
>          return -EBUSY;
>      }
>  
> +    /*
> +     * CPUs except the first online CPU would send a fake (self) NMI to
> +     * rendezvous in NMI handler. But a fake NMI to nmi_cpu may trigger
> +     * unknown_nmi_error(). It ensures nmi_cpu won't receive a fake NMI.
> +     */
> +    ASSERT(cpumask_first(&cpu_online_map) == nmi_cpu);

Looking at this again, I don't think it should be ASSERT(). Instead
you want to return an (easy to recognize) error in this case - maybe
-EPERM or -ENOEXEC or -EXDEV. This is not the least to also be safe
in non-debug builds.

Jan

Chao Gao Sept. 27, 2019, 1:53 p.m. UTC | #2

On Fri, Sep 27, 2019 at 12:19:22PM +0200, Jan Beulich wrote:
>On 26.09.2019 15:53, Chao Gao wrote:
>> @@ -105,23 +110,42 @@ void __init microcode_set_module(unsigned int idx)
>>  }
>>  
>>  /*
>> - * The format is '[<integer>|scan]'. Both options are optional.
>> + * The format is '[<integer>|scan, nmi=<bool>]'. Both options are optional.
>>   * If the EFI has forced which of the multiboot payloads is to be used,
>> - * no parsing will be attempted.
>> + * only nmi=<bool> is parsed.
>>   */
>>  static int __init parse_ucode(const char *s)
>>  {
>> -    const char *q = NULL;
>> +    const char *ss;
>> +    int val, rc = 0;
>>  
>> -    if ( ucode_mod_forced ) /* Forced by EFI */
>> -       return 0;
>> +    do {
>> +        ss = strchr(s, ',');
>> +        if ( !ss )
>> +            ss = strchr(s, '\0');
>>  
>> -    if ( !strncmp(s, "scan", 4) )
>> -        ucode_scan = 1;
>> -    else
>> -        ucode_mod_idx = simple_strtol(s, &q, 0);
>> +        if ( (val = parse_boolean("nmi", s, ss)) >= 0 )
>> +            ucode_in_nmi = val;
>> +        else if ( !ucode_mod_forced ) /* Not forced by EFI */
>> +        {
>> +            const char *q = NULL;
>> +
>> +            if ( !strncmp(s, "scan", 4) )
>> +            {
>> +                ucode_scan = true;
>
>I guess it would have resulted in more consistent code if you had
>used parse_boolean() here, too.
>
>> @@ -222,6 +246,8 @@ const struct microcode_ops *microcode_ops;
>>  static DEFINE_SPINLOCK(microcode_mutex);
>>  
>>  DEFINE_PER_CPU(struct cpu_signature, cpu_sig);
>> +/* Store error code of the work done in NMI handler */
>> +DEFINE_PER_CPU(int, loading_err);
>
>static
>
>> @@ -356,42 +383,88 @@ static void set_state(unsigned int state)
>>      smp_wmb();
>>  }
>>  
>> -static int secondary_thread_fn(void)
>> +static int secondary_nmi_work(void)
>>  {
>> -    unsigned int primary = cpumask_first(this_cpu(cpu_sibling_mask));
>> +    cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
>>  
>> -    if ( !wait_for_state(LOADING_CALLIN) )
>> -        return -EBUSY;
>> +    return wait_for_state(LOADING_EXIT) ? 0 : -EBUSY;
>> +}
>> +
>> +static int primary_thread_work(const struct microcode_patch *patch)
>> +{
>> +    int ret;
>>  
>>      cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
>>  
>> -    if ( !wait_for_state(LOADING_EXIT) )
>> +    if ( !wait_for_state(LOADING_ENTER) )
>>          return -EBUSY;
>>  
>> -    /* Copy update revision from the primary thread. */
>> -    this_cpu(cpu_sig).rev = per_cpu(cpu_sig, primary).rev;
>> +    ret = microcode_ops->apply_microcode(patch);
>> +    if ( !ret )
>> +        atomic_inc(&cpu_updated);
>> +    atomic_inc(&cpu_out);
>>  
>> -    return 0;
>> +    return ret;
>>  }
>>  
>> -static int primary_thread_fn(const struct microcode_patch *patch)
>> +static int primary_nmi_work(const struct microcode_patch *patch)
>> +{
>> +    return primary_thread_work(patch);
>> +}
>
>Why this wrapper? The function signatures are identical. I guess
>you want to emphasize the environment the function is to be used
>in, so perhaps fine despite the redundancy. At least there's no
>address taken of this function, so the compiler can eliminate it.
>
>> +static int secondary_thread_fn(void)
>> +{
>>      if ( !wait_for_state(LOADING_CALLIN) )
>>          return -EBUSY;
>>  
>> -    cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
>> +    self_nmi();
>>  
>> -    if ( !wait_for_state(LOADING_ENTER) )
>> +    /* Copy update revision from the primary thread. */
>> +    this_cpu(cpu_sig).rev =
>> +        per_cpu(cpu_sig, cpumask_first(this_cpu(cpu_sibling_mask))).rev;
>
>_alternative_instructions() takes specific care to avoid relying on
>the NMI potentially not arriving synchronously (in which case you'd
>potentially copy a not-yet-updated CPU signature above). I think the
>same care wants applying here, which I guess would be another
>
>    wait_for_state(LOADING_EXIT);
>
>> +    return this_cpu(loading_err);
>> +}
>> +
>> +static int primary_thread_fn(const struct microcode_patch *patch)
>> +{
>> +    if ( !wait_for_state(LOADING_CALLIN) )
>>          return -EBUSY;
>>  
>> -    ret = microcode_ops->apply_microcode(patch);
>> -    if ( !ret )
>> -        atomic_inc(&cpu_updated);
>> -    atomic_inc(&cpu_out);
>> +    if ( ucode_in_nmi )
>> +    {
>> +        self_nmi();
>> +        return this_cpu(loading_err);
>
>Same here than, to protect against returning a not-yet-updated error
>indicator.
>
>> @@ -420,14 +498,23 @@ static int control_thread_fn(const struct microcode_patch *patch)
>>          return ret;
>>      }
>>  
>> -    /* Let primary threads load the given ucode update */
>> -    set_state(LOADING_ENTER);
>> -
>> +    /* Control thread loads ucode first while others are in NMI handler. */
>>      ret = microcode_ops->apply_microcode(patch);
>>      if ( !ret )
>>          atomic_inc(&cpu_updated);
>>      atomic_inc(&cpu_out);
>>  
>> +    if ( ret == -EIO )
>> +    {
>> +        printk(XENLOG_ERR
>> +               "Late loading aborted: CPU%u failed to update ucode\n", cpu);
>> +        set_state(LOADING_EXIT);
>> +        return ret;
>> +    }
>> +
>> +    /* Let primary threads load the given ucode update */
>> +    set_state(LOADING_ENTER);
>
>While the description goes to some lengths to explain this ordering of
>updates, I still don't really see the point: How is it better for the
>control CPU to have updated its ucode early and then hit an NMI before
>the other CPUs have even started updating, than the other way around
>in the opposite case?

We want to be conservative here. If an ucode is to update something
shared by a whole socket, for the latter case, control thread may
be accessing things that are being updating by the ucode loading on
other cores. It is not safe, just like sibling thread isn't expected
to access features exposed by the old ucode when primary thread is
loading ucode.

Do you think it makes a little sense? If yes, I would like to post
a new version of this patch later this day to catch up Xen 4.13.

Other comments make sense to me.

Thanks
Chao

Jan Beulich Sept. 27, 2019, 1:55 p.m. UTC | #3

On 27.09.2019 15:53, Chao Gao wrote:
> On Fri, Sep 27, 2019 at 12:19:22PM +0200, Jan Beulich wrote:
>> On 26.09.2019 15:53, Chao Gao wrote:
>>> @@ -420,14 +498,23 @@ static int control_thread_fn(const struct microcode_patch *patch)
>>>          return ret;
>>>      }
>>>  
>>> -    /* Let primary threads load the given ucode update */
>>> -    set_state(LOADING_ENTER);
>>> -
>>> +    /* Control thread loads ucode first while others are in NMI handler. */
>>>      ret = microcode_ops->apply_microcode(patch);
>>>      if ( !ret )
>>>          atomic_inc(&cpu_updated);
>>>      atomic_inc(&cpu_out);
>>>  
>>> +    if ( ret == -EIO )
>>> +    {
>>> +        printk(XENLOG_ERR
>>> +               "Late loading aborted: CPU%u failed to update ucode\n", cpu);
>>> +        set_state(LOADING_EXIT);
>>> +        return ret;
>>> +    }
>>> +
>>> +    /* Let primary threads load the given ucode update */
>>> +    set_state(LOADING_ENTER);
>>
>> While the description goes to some lengths to explain this ordering of
>> updates, I still don't really see the point: How is it better for the
>> control CPU to have updated its ucode early and then hit an NMI before
>> the other CPUs have even started updating, than the other way around
>> in the opposite case?
> 
> We want to be conservative here. If an ucode is to update something
> shared by a whole socket, for the latter case, control thread may
> be accessing things that are being updating by the ucode loading on
> other cores. It is not safe, just like sibling thread isn't expected
> to access features exposed by the old ucode when primary thread is
> loading ucode.

Ah yes, considering a socket-wide effect didn't occur to me (although
it should have). So if you mention this aspect in the description, I
think I'm going to be fine with the change in this regard. Yet (as so
often) this raises another question: What about "secondary" sockets?
Shouldn't we entertain a similar two-step approach there then?

Jan

Chao Gao Sept. 27, 2019, 2:15 p.m. UTC | #4

On Fri, Sep 27, 2019 at 03:55:00PM +0200, Jan Beulich wrote:
>On 27.09.2019 15:53, Chao Gao wrote:
>> On Fri, Sep 27, 2019 at 12:19:22PM +0200, Jan Beulich wrote:
>>> On 26.09.2019 15:53, Chao Gao wrote:
>>>> @@ -420,14 +498,23 @@ static int control_thread_fn(const struct microcode_patch *patch)
>>>>          return ret;
>>>>      }
>>>>  
>>>> -    /* Let primary threads load the given ucode update */
>>>> -    set_state(LOADING_ENTER);
>>>> -
>>>> +    /* Control thread loads ucode first while others are in NMI handler. */
>>>>      ret = microcode_ops->apply_microcode(patch);
>>>>      if ( !ret )
>>>>          atomic_inc(&cpu_updated);
>>>>      atomic_inc(&cpu_out);
>>>>  
>>>> +    if ( ret == -EIO )
>>>> +    {
>>>> +        printk(XENLOG_ERR
>>>> +               "Late loading aborted: CPU%u failed to update ucode\n", cpu);
>>>> +        set_state(LOADING_EXIT);
>>>> +        return ret;
>>>> +    }
>>>> +
>>>> +    /* Let primary threads load the given ucode update */
>>>> +    set_state(LOADING_ENTER);
>>>
>>> While the description goes to some lengths to explain this ordering of
>>> updates, I still don't really see the point: How is it better for the
>>> control CPU to have updated its ucode early and then hit an NMI before
>>> the other CPUs have even started updating, than the other way around
>>> in the opposite case?
>> 
>> We want to be conservative here. If an ucode is to update something
>> shared by a whole socket, for the latter case, control thread may
>> be accessing things that are being updating by the ucode loading on
>> other cores. It is not safe, just like sibling thread isn't expected
>> to access features exposed by the old ucode when primary thread is
>> loading ucode.
>
>Ah yes, considering a socket-wide effect didn't occur to me (although
>it should have). So if you mention this aspect in the description, I
>think I'm going to be fine with the change in this regard. Yet (as so
>often) this raises another question: What about "secondary" sockets?
>Shouldn't we entertain a similar two-step approach there then?

No. The two-step approach is because control thread cannot call
self_nmi() in case of triggering unknown_nmi_error() and what is done
in the main NMI handler isn't well controlled. All cores on other
sockets will rendezvous in NMI handler. It means every core's behavior
on other sockets is well controlled.

Thanks
Chao

Patch

diff --git a/docs/misc/xen-command-line.pandoc b/docs/misc/xen-command-line.pandoc
index 832797e..8beb285 100644
--- a/docs/misc/xen-command-line.pandoc
+++ b/docs/misc/xen-command-line.pandoc
@@ -2036,7 +2036,7 @@  pages) must also be specified via the tbuf_size parameter.
 > `= unstable | skewed | stable:socket`
 
 ### ucode (x86)
-> `= [<integer> | scan]`
+> `= List of [ <integer> | scan, nmi=<bool> ]`
 
 Specify how and where to find CPU microcode update blob.
 
@@ -2057,6 +2057,10 @@  microcode in the cpio name space must be:
   - on Intel: kernel/x86/microcode/GenuineIntel.bin
   - on AMD  : kernel/x86/microcode/AuthenticAMD.bin
 
+'nmi' determines late loading is performed in NMI handler or just in
+stop_machine context. In NMI handler, even NMIs are blocked, which is
+considered safer. The default value is `true`.
+
 ### unrestricted_guest (Intel)
 > `= <boolean>`
 
diff --git a/xen/arch/x86/microcode.c b/xen/arch/x86/microcode.c
index 6c23879..b9fa8bb 100644
--- a/xen/arch/x86/microcode.c
+++ b/xen/arch/x86/microcode.c
@@ -36,8 +36,10 @@ 
 #include <xen/earlycpio.h>
 #include <xen/watchdog.h>
 
+#include <asm/apic.h>
 #include <asm/delay.h>
 #include <asm/msr.h>
+#include <asm/nmi.h>
 #include <asm/processor.h>
 #include <asm/setup.h>
 #include <asm/microcode.h>
@@ -95,6 +97,9 @@  static struct ucode_mod_blob __initdata ucode_blob;
  */
 static bool_t __initdata ucode_scan;
 
+/* By default, ucode loading is done in NMI handler */
+static bool ucode_in_nmi = true;
+
 /* Protected by microcode_mutex */
 static struct microcode_patch *microcode_cache;
 
@@ -105,23 +110,42 @@  void __init microcode_set_module(unsigned int idx)
 }
 
 /*
- * The format is '[<integer>|scan]'. Both options are optional.
+ * The format is '[<integer>|scan, nmi=<bool>]'. Both options are optional.
  * If the EFI has forced which of the multiboot payloads is to be used,
- * no parsing will be attempted.
+ * only nmi=<bool> is parsed.
  */
 static int __init parse_ucode(const char *s)
 {
-    const char *q = NULL;
+    const char *ss;
+    int val, rc = 0;
 
-    if ( ucode_mod_forced ) /* Forced by EFI */
-       return 0;
+    do {
+        ss = strchr(s, ',');
+        if ( !ss )
+            ss = strchr(s, '\0');
 
-    if ( !strncmp(s, "scan", 4) )
-        ucode_scan = 1;
-    else
-        ucode_mod_idx = simple_strtol(s, &q, 0);
+        if ( (val = parse_boolean("nmi", s, ss)) >= 0 )
+            ucode_in_nmi = val;
+        else if ( !ucode_mod_forced ) /* Not forced by EFI */
+        {
+            const char *q = NULL;
+
+            if ( !strncmp(s, "scan", 4) )
+            {
+                ucode_scan = true;
+                q = s + 4;
+            }
+            else
+                ucode_mod_idx = simple_strtol(s, &q, 0);
+
+            if ( q != ss )
+                rc = -EINVAL;
+        }
+
+        s = ss + 1;
+    } while ( *ss );
 
-    return (q && *q) ? -EINVAL : 0;
+    return rc;
 }
 custom_param("ucode", parse_ucode);
 
@@ -222,6 +246,8 @@  const struct microcode_ops *microcode_ops;
 static DEFINE_SPINLOCK(microcode_mutex);
 
 DEFINE_PER_CPU(struct cpu_signature, cpu_sig);
+/* Store error code of the work done in NMI handler */
+DEFINE_PER_CPU(int, loading_err);
 
 /*
  * Count the CPUs that have entered, exited the rendezvous and succeeded in
@@ -232,6 +258,7 @@  DEFINE_PER_CPU(struct cpu_signature, cpu_sig);
  */
 static cpumask_t cpu_callin_map;
 static atomic_t cpu_out, cpu_updated;
+static const struct microcode_patch *nmi_patch = ZERO_BLOCK_PTR;
 
 /*
  * Return a patch that covers current CPU. If there are multiple patches,
@@ -356,42 +383,88 @@  static void set_state(unsigned int state)
     smp_wmb();
 }
 
-static int secondary_thread_fn(void)
+static int secondary_nmi_work(void)
 {
-    unsigned int primary = cpumask_first(this_cpu(cpu_sibling_mask));
+    cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
 
-    if ( !wait_for_state(LOADING_CALLIN) )
-        return -EBUSY;
+    return wait_for_state(LOADING_EXIT) ? 0 : -EBUSY;
+}
+
+static int primary_thread_work(const struct microcode_patch *patch)
+{
+    int ret;
 
     cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
 
-    if ( !wait_for_state(LOADING_EXIT) )
+    if ( !wait_for_state(LOADING_ENTER) )
         return -EBUSY;
 
-    /* Copy update revision from the primary thread. */
-    this_cpu(cpu_sig).rev = per_cpu(cpu_sig, primary).rev;
+    ret = microcode_ops->apply_microcode(patch);
+    if ( !ret )
+        atomic_inc(&cpu_updated);
+    atomic_inc(&cpu_out);
 
-    return 0;
+    return ret;
 }
 
-static int primary_thread_fn(const struct microcode_patch *patch)
+static int primary_nmi_work(const struct microcode_patch *patch)
+{
+    return primary_thread_work(patch);
+}
+
+static int microcode_nmi_callback(const struct cpu_user_regs *regs, int cpu)
 {
-    int ret = 0;
+    unsigned int primary = cpumask_first(this_cpu(cpu_sibling_mask));
+    int ret;
+
+    /* System-generated NMI, leave to main handler */
+    if ( ACCESS_ONCE(loading_state) != LOADING_CALLIN )
+        return 0;
+
+    /*
+     * Primary threads load ucode in NMI handler on if ucode_in_nmi is true.
+     * Secondary threads are expected to stay in NMI handler regardless of
+     * ucode_in_nmi.
+     */
+    if ( cpu == cpumask_first(&cpu_online_map) ||
+         (!ucode_in_nmi && cpu == primary) )
+        return 0;
+
+    if ( cpu == primary )
+        ret = primary_nmi_work(nmi_patch);
+    else
+        ret = secondary_nmi_work();
+    this_cpu(loading_err) = ret;
+
+    return 0;
+}
 
+static int secondary_thread_fn(void)
+{
     if ( !wait_for_state(LOADING_CALLIN) )
         return -EBUSY;
 
-    cpumask_set_cpu(smp_processor_id(), &cpu_callin_map);
+    self_nmi();
 
-    if ( !wait_for_state(LOADING_ENTER) )
+    /* Copy update revision from the primary thread. */
+    this_cpu(cpu_sig).rev =
+        per_cpu(cpu_sig, cpumask_first(this_cpu(cpu_sibling_mask))).rev;
+
+    return this_cpu(loading_err);
+}
+
+static int primary_thread_fn(const struct microcode_patch *patch)
+{
+    if ( !wait_for_state(LOADING_CALLIN) )
         return -EBUSY;
 
-    ret = microcode_ops->apply_microcode(patch);
-    if ( !ret )
-        atomic_inc(&cpu_updated);
-    atomic_inc(&cpu_out);
+    if ( ucode_in_nmi )
+    {
+        self_nmi();
+        return this_cpu(loading_err);
+    }
 
-    return ret;
+    return primary_thread_work(patch);
 }
 
 static int control_thread_fn(const struct microcode_patch *patch)
@@ -399,6 +472,7 @@  static int control_thread_fn(const struct microcode_patch *patch)
     unsigned int cpu = smp_processor_id(), done;
     unsigned long tick;
     int ret;
+    nmi_callback_t *saved_nmi_callback;
 
     /*
      * We intend to keep interrupt disabled for a long time, which may lead to
@@ -406,6 +480,10 @@  static int control_thread_fn(const struct microcode_patch *patch)
      */
     watchdog_disable();
 
+    nmi_patch = patch;
+    smp_wmb();
+    saved_nmi_callback = set_nmi_callback(microcode_nmi_callback);
+
     /* Allow threads to call in */
     set_state(LOADING_CALLIN);
 
@@ -420,14 +498,23 @@  static int control_thread_fn(const struct microcode_patch *patch)
         return ret;
     }
 
-    /* Let primary threads load the given ucode update */
-    set_state(LOADING_ENTER);
-
+    /* Control thread loads ucode first while others are in NMI handler. */
     ret = microcode_ops->apply_microcode(patch);
     if ( !ret )
         atomic_inc(&cpu_updated);
     atomic_inc(&cpu_out);
 
+    if ( ret == -EIO )
+    {
+        printk(XENLOG_ERR
+               "Late loading aborted: CPU%u failed to update ucode\n", cpu);
+        set_state(LOADING_EXIT);
+        return ret;
+    }
+
+    /* Let primary threads load the given ucode update */
+    set_state(LOADING_ENTER);
+
     tick = rdtsc_ordered();
     /* Wait for primary threads finishing update */
     while ( (done = atomic_read(&cpu_out)) != nr_cores )
@@ -456,6 +543,8 @@  static int control_thread_fn(const struct microcode_patch *patch)
     /* Mark loading is done to unblock other threads */
     set_state(LOADING_EXIT);
 
+    set_nmi_callback(saved_nmi_callback);
+    nmi_patch = ZERO_BLOCK_PTR;
     watchdog_enable();
 
     return ret;
@@ -515,6 +604,13 @@  int microcode_update(XEN_GUEST_HANDLE_PARAM(const_void) buf, unsigned long len)
         return -EBUSY;
     }
 
+    /*
+     * CPUs except the first online CPU would send a fake (self) NMI to
+     * rendezvous in NMI handler. But a fake NMI to nmi_cpu may trigger
+     * unknown_nmi_error(). It ensures nmi_cpu won't receive a fake NMI.
+     */
+    ASSERT(cpumask_first(&cpu_online_map) == nmi_cpu);
+
     patch = parse_blob(buffer, len);
     xfree(buffer);
     if ( IS_ERR(patch) )
diff --git a/xen/arch/x86/traps.c b/xen/arch/x86/traps.c
index 16c590d..2cd5e29 100644
--- a/xen/arch/x86/traps.c
+++ b/xen/arch/x86/traps.c
@@ -126,6 +126,8 @@  boolean_param("ler", opt_ler);
 /* LastExceptionFromIP on this hardware.  Zero if LER is not in use. */
 unsigned int __read_mostly ler_msr;
 
+const unsigned int nmi_cpu;
+
 #define stack_words_per_line 4
 #define ESP_BEFORE_EXCEPTION(regs) ((unsigned long *)regs->rsp)
 
@@ -1679,7 +1681,7 @@  void do_nmi(const struct cpu_user_regs *regs)
      * this port before we re-arm the NMI watchdog, we reduce the chance
      * of having an NMI watchdog expire while in the SMI handler.
      */
-    if ( cpu == 0 )
+    if ( cpu == nmi_cpu )
         reason = inb(0x61);
 
     if ( (nmi_watchdog == NMI_NONE) ||
@@ -1687,7 +1689,7 @@  void do_nmi(const struct cpu_user_regs *regs)
         handle_unknown = true;
 
     /* Only the BSP gets external NMIs from the system. */
-    if ( cpu == 0 )
+    if ( cpu == nmi_cpu )
     {
         if ( reason & 0x80 )
             pci_serr_error(regs);
diff --git a/xen/include/asm-x86/nmi.h b/xen/include/asm-x86/nmi.h
index 99f6284..f9dfca6 100644
--- a/xen/include/asm-x86/nmi.h
+++ b/xen/include/asm-x86/nmi.h
@@ -11,6 +11,9 @@  extern bool opt_watchdog;
 
 /* Watchdog force parameter from the command line */
 extern bool watchdog_force;
+
+/* CPU to handle platform NMI */
+extern const unsigned int nmi_cpu;
  
 typedef int nmi_callback_t(const struct cpu_user_regs *regs, int cpu);