]> git.proxmox.com Git - mirror_edk2.git/blobdiff - UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
UefiCpuPkg/PiSmmCpuDxeSmm: Make code consistent with comments
[mirror_edk2.git] / UefiCpuPkg / PiSmmCpuDxeSmm / MpService.c
index 42b040531e7f497991271427a5bd6cef502d73d2..328d7ae53a492e851d426dce1970424b551a9d97 100644 (file)
@@ -1,16 +1,10 @@
 /** @file\r
 SMM MP service implementation\r
 \r
-Copyright (c) 2009 - 2018, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2009 - 2019, Intel Corporation. All rights reserved.<BR>\r
 Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>\r
 \r
-This program and the accompanying materials\r
-are licensed and made available under the terms and conditions of the BSD License\r
-which accompanies this distribution.  The full text of the license may be found at\r
-http://opensource.org/licenses/bsd-license.php\r
-\r
-THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
-WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
@@ -142,11 +136,9 @@ ReleaseAllAPs (
   )\r
 {\r
   UINTN                             Index;\r
-  UINTN                             BspIndex;\r
 \r
-  BspIndex = mSmmMpSyncData->BspIndex;\r
   for (Index = mMaxNumberOfCpus; Index-- > 0;) {\r
-    if (Index != BspIndex && *(mSmmMpSyncData->CpuData[Index].Present)) {\r
+    if (IsPresentAp (Index)) {\r
       ReleaseSemaphore (mSmmMpSyncData->CpuData[Index].Run);\r
     }\r
   }\r
@@ -353,6 +345,165 @@ ReplaceOSMtrrs (
   MtrrSetAllMtrrs (&gSmiMtrrs);\r
 }\r
 \r
+/**\r
+  Wheck whether task has been finished by all APs.\r
+\r
+  @param       BlockMode   Whether did it in block mode or non-block mode.\r
+\r
+  @retval      TRUE        Task has been finished by all APs.\r
+  @retval      FALSE       Task not has been finished by all APs.\r
+\r
+**/\r
+BOOLEAN\r
+WaitForAllAPsNotBusy (\r
+  IN BOOLEAN                        BlockMode\r
+  )\r
+{\r
+  UINTN                             Index;\r
+\r
+  for (Index = mMaxNumberOfCpus; Index-- > 0;) {\r
+    //\r
+    // Ignore BSP and APs which not call in SMM.\r
+    //\r
+    if (!IsPresentAp(Index)) {\r
+      continue;\r
+    }\r
+\r
+    if (BlockMode) {\r
+      AcquireSpinLock(mSmmMpSyncData->CpuData[Index].Busy);\r
+      ReleaseSpinLock(mSmmMpSyncData->CpuData[Index].Busy);\r
+    } else {\r
+      if (AcquireSpinLockOrFail (mSmmMpSyncData->CpuData[Index].Busy)) {\r
+        ReleaseSpinLock(mSmmMpSyncData->CpuData[Index].Busy);\r
+      } else {\r
+        return FALSE;\r
+      }\r
+    }\r
+  }\r
+\r
+  return TRUE;\r
+}\r
+\r
+/**\r
+  Check whether it is an present AP.\r
+\r
+  @param   CpuIndex      The AP index which calls this function.\r
+\r
+  @retval  TRUE           It's a present AP.\r
+  @retval  TRUE           This is not an AP or it is not present.\r
+\r
+**/\r
+BOOLEAN\r
+IsPresentAp (\r
+  IN UINTN        CpuIndex\r
+  )\r
+{\r
+  return ((CpuIndex != gSmmCpuPrivate->SmmCoreEntryContext.CurrentlyExecutingCpu) &&\r
+    *(mSmmMpSyncData->CpuData[CpuIndex].Present));\r
+}\r
+\r
+/**\r
+  Check whether execute in single AP or all APs.\r
+\r
+  Compare two Tokens used by different APs to know whether in StartAllAps call.\r
+\r
+  Whether is an valid AP base on AP's Present flag.\r
+\r
+  @retval  TRUE      IN StartAllAps call.\r
+  @retval  FALSE     Not in StartAllAps call.\r
+\r
+**/\r
+BOOLEAN\r
+InStartAllApsCall (\r
+  VOID\r
+  )\r
+{\r
+  UINTN      ApIndex;\r
+  UINTN      ApIndex2;\r
+\r
+  for (ApIndex = mMaxNumberOfCpus; ApIndex-- > 0;) {\r
+    if (IsPresentAp (ApIndex) && (mSmmMpSyncData->CpuData[ApIndex].Token != NULL)) {\r
+      for (ApIndex2 = ApIndex; ApIndex2-- > 0;) {\r
+        if (IsPresentAp (ApIndex2) && (mSmmMpSyncData->CpuData[ApIndex2].Token != NULL)) {\r
+          return mSmmMpSyncData->CpuData[ApIndex2].Token == mSmmMpSyncData->CpuData[ApIndex].Token;\r
+        }\r
+      }\r
+    }\r
+  }\r
+\r
+  return FALSE;\r
+}\r
+\r
+/**\r
+  Clean up the status flags used during executing the procedure.\r
+\r
+  @param   CpuIndex      The AP index which calls this function.\r
+\r
+**/\r
+VOID\r
+ReleaseToken (\r
+  IN UINTN                  CpuIndex\r
+  )\r
+{\r
+  UINTN                             Index;\r
+  BOOLEAN                           Released;\r
+\r
+  if (InStartAllApsCall ()) {\r
+    //\r
+    // In Start All APs mode, make sure all APs have finished task.\r
+    //\r
+    if (WaitForAllAPsNotBusy (FALSE)) {\r
+      //\r
+      // Clean the flags update in the function call.\r
+      //\r
+      Released = FALSE;\r
+      for (Index = mMaxNumberOfCpus; Index-- > 0;) {\r
+        //\r
+        // Only In SMM APs need to be clean up.\r
+        //\r
+        if (mSmmMpSyncData->CpuData[Index].Present && mSmmMpSyncData->CpuData[Index].Token != NULL) {\r
+          if (!Released) {\r
+            ReleaseSpinLock (mSmmMpSyncData->CpuData[Index].Token);\r
+            Released = TRUE;\r
+          }\r
+          mSmmMpSyncData->CpuData[Index].Token = NULL;\r
+        }\r
+      }\r
+    }\r
+  } else {\r
+    //\r
+    // In single AP mode.\r
+    //\r
+    if (mSmmMpSyncData->CpuData[CpuIndex].Token != NULL) {\r
+      ReleaseSpinLock (mSmmMpSyncData->CpuData[CpuIndex].Token);\r
+      mSmmMpSyncData->CpuData[CpuIndex].Token = NULL;\r
+    }\r
+  }\r
+}\r
+\r
+/**\r
+  Free the tokens in the maintained list.\r
+\r
+**/\r
+VOID\r
+FreeTokens (\r
+  VOID\r
+  )\r
+{\r
+  LIST_ENTRY            *Link;\r
+  PROCEDURE_TOKEN       *ProcToken;\r
+\r
+  while (!IsListEmpty (&gSmmCpuPrivate->TokenList)) {\r
+    Link = GetFirstNode (&gSmmCpuPrivate->TokenList);\r
+    ProcToken = PROCEDURE_TOKEN_FROM_LINK (Link);\r
+\r
+    RemoveEntryList (&ProcToken->Link);\r
+\r
+    FreePool ((VOID *)ProcToken->ProcedureToken);\r
+    FreePool (ProcToken);\r
+  }\r
+}\r
+\r
 /**\r
   SMI handler for BSP.\r
 \r
@@ -482,12 +633,7 @@ BSPHandler (
   //\r
   // Make sure all APs have completed their pending none-block tasks\r
   //\r
-  for (Index = mMaxNumberOfCpus; Index-- > 0;) {\r
-    if (Index != CpuIndex && *(mSmmMpSyncData->CpuData[Index].Present)) {\r
-      AcquireSpinLock (mSmmMpSyncData->CpuData[Index].Busy);\r
-      ReleaseSpinLock (mSmmMpSyncData->CpuData[Index].Busy);\r
-    }\r
-  }\r
+  WaitForAllAPsNotBusy (TRUE);\r
 \r
   //\r
   // Perform the remaining tasks\r
@@ -578,6 +724,11 @@ BSPHandler (
   //\r
   WaitForAllAPs (ApCount);\r
 \r
+  //\r
+  // Clean the tokens buffer.\r
+  //\r
+  FreeTokens ();\r
+\r
   //\r
   // Reset BspIndex to -1, meaning BSP has not been elected.\r
   //\r
@@ -610,6 +761,7 @@ APHandler (
   UINT64                            Timer;\r
   UINTN                             BspIndex;\r
   MTRR_SETTINGS                     Mtrrs;\r
+  EFI_STATUS                        ProcedureStatus;\r
 \r
   //\r
   // Timeout BSP\r
@@ -736,14 +888,19 @@ APHandler (
     //\r
     // Invoke the scheduled procedure\r
     //\r
-    (*mSmmMpSyncData->CpuData[CpuIndex].Procedure) (\r
-      (VOID*)mSmmMpSyncData->CpuData[CpuIndex].Parameter\r
-      );\r
+    ProcedureStatus = (*mSmmMpSyncData->CpuData[CpuIndex].Procedure) (\r
+                          (VOID*)mSmmMpSyncData->CpuData[CpuIndex].Parameter\r
+                          );\r
+    if (mSmmMpSyncData->CpuData[CpuIndex].Status != NULL) {\r
+      *mSmmMpSyncData->CpuData[CpuIndex].Status = ProcedureStatus;\r
+    }\r
 \r
     //\r
     // Release BUSY\r
     //\r
     ReleaseSpinLock (mSmmMpSyncData->CpuData[CpuIndex].Busy);\r
+\r
+    ReleaseToken (CpuIndex);\r
   }\r
 \r
   if (SmmCpuFeaturesNeedConfigureMtrrs()) {\r
@@ -912,13 +1069,124 @@ Gen4GPageTable (
   return (UINT32)(UINTN)PageTable;\r
 }\r
 \r
+/**\r
+  Checks whether the input token is the current used token.\r
+\r
+  @param[in]  Token      This parameter describes the token that was passed into DispatchProcedure or\r
+                         BroadcastProcedure.\r
+\r
+  @retval TRUE           The input token is the current used token.\r
+  @retval FALSE          The input token is not the current used token.\r
+**/\r
+BOOLEAN\r
+IsTokenInUse (\r
+  IN SPIN_LOCK           *Token\r
+  )\r
+{\r
+  LIST_ENTRY        *Link;\r
+  PROCEDURE_TOKEN   *ProcToken;\r
+\r
+  if (Token == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  Link = GetFirstNode (&gSmmCpuPrivate->TokenList);\r
+  while (!IsNull (&gSmmCpuPrivate->TokenList, Link)) {\r
+    ProcToken = PROCEDURE_TOKEN_FROM_LINK (Link);\r
+\r
+    if (ProcToken->ProcedureToken == Token) {\r
+      return TRUE;\r
+    }\r
+\r
+    Link = GetNextNode (&gSmmCpuPrivate->TokenList, Link);\r
+  }\r
+\r
+  return FALSE;\r
+}\r
+\r
+/**\r
+  create token and save it to the maintain list.\r
+\r
+  @retval    return the spin lock used as token.\r
+\r
+**/\r
+SPIN_LOCK *\r
+CreateToken (\r
+  VOID\r
+  )\r
+{\r
+  PROCEDURE_TOKEN    *ProcToken;\r
+  SPIN_LOCK           *CpuToken;\r
+  UINTN               SpinLockSize;\r
+\r
+  SpinLockSize = GetSpinLockProperties ();\r
+  CpuToken = AllocatePool (SpinLockSize);\r
+  ASSERT (CpuToken != NULL);\r
+  InitializeSpinLock (CpuToken);\r
+  AcquireSpinLock (CpuToken);\r
+\r
+  ProcToken = AllocatePool (sizeof (PROCEDURE_TOKEN));\r
+  ASSERT (ProcToken != NULL);\r
+  ProcToken->Signature = PROCEDURE_TOKEN_SIGNATURE;\r
+  ProcToken->ProcedureToken = CpuToken;\r
+\r
+  InsertTailList (&gSmmCpuPrivate->TokenList, &ProcToken->Link);\r
+\r
+  return CpuToken;\r
+}\r
+\r
+/**\r
+  Checks status of specified AP.\r
+\r
+  This function checks whether the specified AP has finished the task assigned\r
+  by StartupThisAP(), and whether timeout expires.\r
+\r
+  @param[in]  Token             This parameter describes the token that was passed into DispatchProcedure or\r
+                                BroadcastProcedure.\r
+\r
+  @retval EFI_SUCCESS           Specified AP has finished task assigned by StartupThisAPs().\r
+  @retval EFI_NOT_READY         Specified AP has not finished task and timeout has not expired.\r
+**/\r
+EFI_STATUS\r
+IsApReady (\r
+  IN SPIN_LOCK          *Token\r
+  )\r
+{\r
+  if (AcquireSpinLockOrFail (Token)) {\r
+    ReleaseSpinLock (Token);\r
+    return EFI_SUCCESS;\r
+  }\r
+\r
+  return EFI_NOT_READY;\r
+}\r
+\r
 /**\r
   Schedule a procedure to run on the specified CPU.\r
 \r
   @param[in]       Procedure                The address of the procedure to run\r
   @param[in]       CpuIndex                 Target CPU Index\r
-  @param[in, out]  ProcArguments            The parameter to pass to the procedure\r
-  @param[in]       BlockingMode             Startup AP in blocking mode or not\r
+  @param[in,out]   ProcArguments            The parameter to pass to the procedure\r
+  @param[in]       Token                    This is an optional parameter that allows the caller to execute the\r
+                                            procedure in a blocking or non-blocking fashion. If it is NULL the\r
+                                            call is blocking, and the call will not return until the AP has\r
+                                            completed the procedure. If the token is not NULL, the call will\r
+                                            return immediately. The caller can check whether the procedure has\r
+                                            completed with CheckOnProcedure or WaitForProcedure.\r
+  @param[in]       TimeoutInMicroseconds    Indicates the time limit in microseconds for the APs to finish\r
+                                            execution of Procedure, either for blocking or non-blocking mode.\r
+                                            Zero means infinity. If the timeout expires before all APs return\r
+                                            from Procedure, then Procedure on the failed APs is terminated. If\r
+                                            the timeout expires in blocking mode, the call returns EFI_TIMEOUT.\r
+                                            If the timeout expires in non-blocking mode, the timeout determined\r
+                                            can be through CheckOnProcedure or WaitForProcedure.\r
+                                            Note that timeout support is optional. Whether an implementation\r
+                                            supports this feature can be determined via the Attributes data\r
+                                            member.\r
+  @param[in,out]   CpuStatus                This optional pointer may be used to get the status code returned\r
+                                            by Procedure when it completes execution on the target AP, or with\r
+                                            EFI_TIMEOUT if the Procedure fails to complete within the optional\r
+                                            timeout. The implementation will update this variable with\r
+                                            EFI_NOT_READY prior to starting Procedure on the target AP.\r
 \r
   @retval EFI_INVALID_PARAMETER    CpuNumber not valid\r
   @retval EFI_INVALID_PARAMETER    CpuNumber specifying BSP\r
@@ -929,10 +1197,12 @@ Gen4GPageTable (
 **/\r
 EFI_STATUS\r
 InternalSmmStartupThisAp (\r
-  IN      EFI_AP_PROCEDURE          Procedure,\r
-  IN      UINTN                     CpuIndex,\r
-  IN OUT  VOID                      *ProcArguments OPTIONAL,\r
-  IN      BOOLEAN                   BlockingMode\r
+  IN      EFI_AP_PROCEDURE2              Procedure,\r
+  IN      UINTN                          CpuIndex,\r
+  IN OUT  VOID                           *ProcArguments OPTIONAL,\r
+  IN      MM_COMPLETION                  *Token,\r
+  IN      UINTN                          TimeoutInMicroseconds,\r
+  IN OUT  EFI_STATUS                     *CpuStatus\r
   )\r
 {\r
   if (CpuIndex >= gSmmCpuPrivate->SmmCoreEntryContext.NumberOfCpus) {\r
@@ -958,24 +1228,190 @@ InternalSmmStartupThisAp (
     }\r
     return EFI_INVALID_PARAMETER;\r
   }\r
+  if ((TimeoutInMicroseconds != 0) && ((mSmmMp.Attributes & EFI_MM_MP_TIMEOUT_SUPPORTED) == 0)) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+  if (Procedure == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
 \r
-  if (BlockingMode) {\r
+  if (Token == NULL) {\r
     AcquireSpinLock (mSmmMpSyncData->CpuData[CpuIndex].Busy);\r
   } else {\r
     if (!AcquireSpinLockOrFail (mSmmMpSyncData->CpuData[CpuIndex].Busy)) {\r
-      DEBUG((DEBUG_ERROR, "mSmmMpSyncData->CpuData[%d].Busy\n", CpuIndex));\r
-      return EFI_INVALID_PARAMETER;\r
+      DEBUG((DEBUG_ERROR, "Can't acquire mSmmMpSyncData->CpuData[%d].Busy\n", CpuIndex));\r
+      return EFI_NOT_READY;\r
     }\r
+\r
+    *Token = (MM_COMPLETION) CreateToken ();\r
   }\r
 \r
   mSmmMpSyncData->CpuData[CpuIndex].Procedure = Procedure;\r
   mSmmMpSyncData->CpuData[CpuIndex].Parameter = ProcArguments;\r
+  if (Token != NULL) {\r
+    mSmmMpSyncData->CpuData[CpuIndex].Token   = (SPIN_LOCK *)(*Token);\r
+  }\r
+  mSmmMpSyncData->CpuData[CpuIndex].Status    = CpuStatus;\r
+  if (mSmmMpSyncData->CpuData[CpuIndex].Status != NULL) {\r
+    *mSmmMpSyncData->CpuData[CpuIndex].Status = EFI_NOT_READY;\r
+  }\r
+\r
   ReleaseSemaphore (mSmmMpSyncData->CpuData[CpuIndex].Run);\r
 \r
-  if (BlockingMode) {\r
+  if (Token == NULL) {\r
     AcquireSpinLock (mSmmMpSyncData->CpuData[CpuIndex].Busy);\r
     ReleaseSpinLock (mSmmMpSyncData->CpuData[CpuIndex].Busy);\r
   }\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Worker function to execute a caller provided function on all enabled APs.\r
+\r
+  @param[in]     Procedure               A pointer to the function to be run on\r
+                                         enabled APs of the system.\r
+  @param[in]     TimeoutInMicroseconds   Indicates the time limit in microseconds for\r
+                                         APs to return from Procedure, either for\r
+                                         blocking or non-blocking mode.\r
+  @param[in,out] ProcedureArguments      The parameter passed into Procedure for\r
+                                         all APs.\r
+  @param[in,out] Token                   This is an optional parameter that allows the caller to execute the\r
+                                         procedure in a blocking or non-blocking fashion. If it is NULL the\r
+                                         call is blocking, and the call will not return until the AP has\r
+                                         completed the procedure. If the token is not NULL, the call will\r
+                                         return immediately. The caller can check whether the procedure has\r
+                                         completed with CheckOnProcedure or WaitForProcedure.\r
+  @param[in,out] CPUStatus               This optional pointer may be used to get the status code returned\r
+                                         by Procedure when it completes execution on the target AP, or with\r
+                                         EFI_TIMEOUT if the Procedure fails to complete within the optional\r
+                                         timeout. The implementation will update this variable with\r
+                                         EFI_NOT_READY prior to starting Procedure on the target AP.\r
+\r
+\r
+  @retval EFI_SUCCESS             In blocking mode, all APs have finished before\r
+                                  the timeout expired.\r
+  @retval EFI_SUCCESS             In non-blocking mode, function has been dispatched\r
+                                  to all enabled APs.\r
+  @retval others                  Failed to Startup all APs.\r
+\r
+**/\r
+EFI_STATUS\r
+InternalSmmStartupAllAPs (\r
+  IN       EFI_AP_PROCEDURE2             Procedure,\r
+  IN       UINTN                         TimeoutInMicroseconds,\r
+  IN OUT   VOID                          *ProcedureArguments OPTIONAL,\r
+  IN OUT   MM_COMPLETION                 *Token,\r
+  IN OUT   EFI_STATUS                    *CPUStatus\r
+  )\r
+{\r
+  UINTN               Index;\r
+  UINTN               CpuCount;\r
+\r
+  if ((TimeoutInMicroseconds != 0) && ((mSmmMp.Attributes & EFI_MM_MP_TIMEOUT_SUPPORTED) == 0)) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+  if (Procedure == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  CpuCount = 0;\r
+  for (Index = mMaxNumberOfCpus; Index-- > 0;) {\r
+    if (IsPresentAp (Index)) {\r
+      CpuCount ++;\r
+\r
+      if (gSmmCpuPrivate->Operation[Index] == SmmCpuRemove) {\r
+        return EFI_INVALID_PARAMETER;\r
+      }\r
+\r
+      if (!AcquireSpinLockOrFail(mSmmMpSyncData->CpuData[Index].Busy)) {\r
+        return EFI_NOT_READY;\r
+      }\r
+      ReleaseSpinLock (mSmmMpSyncData->CpuData[Index].Busy);\r
+    }\r
+  }\r
+  if (CpuCount == 0) {\r
+    return EFI_NOT_STARTED;\r
+  }\r
+\r
+  if (Token != NULL) {\r
+    *Token = (MM_COMPLETION) CreateToken ();\r
+  }\r
+\r
+  //\r
+  // Make sure all BUSY should be acquired.\r
+  //\r
+  // Because former code already check mSmmMpSyncData->CpuData[***].Busy for each AP.\r
+  // Here code always use AcquireSpinLock instead of AcquireSpinLockOrFail for not\r
+  // block mode.\r
+  //\r
+  for (Index = mMaxNumberOfCpus; Index-- > 0;) {\r
+    if (IsPresentAp (Index)) {\r
+      AcquireSpinLock (mSmmMpSyncData->CpuData[Index].Busy);\r
+    }\r
+  }\r
+\r
+  for (Index = mMaxNumberOfCpus; Index-- > 0;) {\r
+    if (IsPresentAp (Index)) {\r
+      mSmmMpSyncData->CpuData[Index].Procedure = (EFI_AP_PROCEDURE2) Procedure;\r
+      mSmmMpSyncData->CpuData[Index].Parameter = ProcedureArguments;\r
+      if (Token != NULL) {\r
+        mSmmMpSyncData->CpuData[Index].Token   = (SPIN_LOCK *)(*Token);\r
+      }\r
+      if (CPUStatus != NULL) {\r
+        mSmmMpSyncData->CpuData[Index].Status    = &CPUStatus[Index];\r
+        if (mSmmMpSyncData->CpuData[Index].Status != NULL) {\r
+          *mSmmMpSyncData->CpuData[Index].Status = EFI_NOT_READY;\r
+        }\r
+      }\r
+    } else {\r
+      //\r
+      // PI spec requirement:\r
+      // For every excluded processor, the array entry must contain a value of EFI_NOT_STARTED.\r
+      //\r
+      if (CPUStatus != NULL) {\r
+        CPUStatus[Index] = EFI_NOT_STARTED;\r
+      }\r
+    }\r
+  }\r
+\r
+  ReleaseAllAPs ();\r
+\r
+  if (Token == NULL) {\r
+    //\r
+    // Make sure all APs have completed their tasks.\r
+    //\r
+    WaitForAllAPsNotBusy (TRUE);\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  ISO C99 6.5.2.2 "Function calls", paragraph 9:\r
+  If the function is defined with a type that is not compatible with\r
+  the type (of the expression) pointed to by the expression that\r
+  denotes the called function, the behavior is undefined.\r
+\r
+  So add below wrapper function to convert between EFI_AP_PROCEDURE\r
+  and EFI_AP_PROCEDURE2.\r
+\r
+  Wrapper for Procedures.\r
+\r
+  @param[in]  Buffer              Pointer to PROCEDURE_WRAPPER buffer.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProcedureWrapper (\r
+  IN     VOID *Buffer\r
+  )\r
+{\r
+  PROCEDURE_WRAPPER *Wrapper;\r
+\r
+  Wrapper = Buffer;\r
+  Wrapper->Procedure (Wrapper->ProcedureArgument);\r
+\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -1001,7 +1437,15 @@ SmmBlockingStartupThisAp (
   IN OUT  VOID                      *ProcArguments OPTIONAL\r
   )\r
 {\r
-  return InternalSmmStartupThisAp(Procedure, CpuIndex, ProcArguments, TRUE);\r
+  PROCEDURE_WRAPPER  Wrapper;\r
+\r
+  Wrapper.Procedure = Procedure;\r
+  Wrapper.ProcedureArgument = ProcArguments;\r
+\r
+  //\r
+  // Use wrapper function to convert EFI_AP_PROCEDURE to EFI_AP_PROCEDURE2.\r
+  //\r
+  return InternalSmmStartupThisAp (ProcedureWrapper, CpuIndex, &Wrapper, NULL, 0, NULL);\r
 }\r
 \r
 /**\r
@@ -1026,7 +1470,22 @@ SmmStartupThisAp (
   IN OUT  VOID                      *ProcArguments OPTIONAL\r
   )\r
 {\r
-  return InternalSmmStartupThisAp(Procedure, CpuIndex, ProcArguments, FeaturePcdGet (PcdCpuSmmBlockStartupThisAp));\r
+  MM_COMPLETION               Token;\r
+\r
+  gSmmCpuPrivate->ApWrapperFunc[CpuIndex].Procedure = Procedure;\r
+  gSmmCpuPrivate->ApWrapperFunc[CpuIndex].ProcedureArgument = ProcArguments;\r
+\r
+  //\r
+  // Use wrapper function to convert EFI_AP_PROCEDURE to EFI_AP_PROCEDURE2.\r
+  //\r
+  return InternalSmmStartupThisAp (\r
+    ProcedureWrapper,\r
+    CpuIndex,\r
+    &gSmmCpuPrivate->ApWrapperFunc[CpuIndex],\r
+    FeaturePcdGet (PcdCpuSmmBlockStartupThisAp) ? NULL : &Token,\r
+    0,\r
+    NULL\r
+    );\r
 }\r
 \r
 /**\r
@@ -1112,9 +1571,18 @@ SmiRendezvous (
   ASSERT(CpuIndex < mMaxNumberOfCpus);\r
 \r
   //\r
-  // Save Cr2 because Page Fault exception in SMM may override its value\r
+  // Save Cr2 because Page Fault exception in SMM may override its value,\r
+  // when using on-demand paging for above 4G memory.\r
   //\r
-  Cr2 = AsmReadCr2 ();\r
+  Cr2 = 0;\r
+  SaveCr2 (&Cr2);\r
+\r
+  //\r
+  // Call the user register Startup function first.\r
+  //\r
+  if (mSmmMpSyncData->StartupProcedure != NULL) {\r
+    mSmmMpSyncData->StartupProcedure (mSmmMpSyncData->StartupProcArgs);\r
+  }\r
 \r
   //\r
   // Perform CPU specific entry hooks\r
@@ -1253,10 +1721,26 @@ SmiRendezvous (
 \r
 Exit:\r
   SmmCpuFeaturesRendezvousExit (CpuIndex);\r
+\r
   //\r
   // Restore Cr2\r
   //\r
-  AsmWriteCr2 (Cr2);\r
+  RestoreCr2 (Cr2);\r
+}\r
+\r
+/**\r
+  Allocate buffer for SpinLock and Wrapper function buffer.\r
+\r
+**/\r
+VOID\r
+InitializeDataForMmMp (\r
+  VOID\r
+  )\r
+{\r
+  gSmmCpuPrivate->ApWrapperFunc = AllocatePool (sizeof (PROCEDURE_WRAPPER) * gSmmCpuPrivate->SmmCoreEntryContext.NumberOfCpus);\r
+  ASSERT (gSmmCpuPrivate->ApWrapperFunc != NULL);\r
+\r
+  InitializeListHead (&gSmmCpuPrivate->TokenList);\r
 }\r
 \r
 /**\r
@@ -1272,7 +1756,6 @@ InitializeSmmCpuSemaphores (
   UINTN                      TotalSize;\r
   UINTN                      GlobalSemaphoresSize;\r
   UINTN                      CpuSemaphoresSize;\r
-  UINTN                      MsrSemahporeSize;\r
   UINTN                      SemaphoreSize;\r
   UINTN                      Pages;\r
   UINTN                      *SemaphoreBlock;\r
@@ -1282,8 +1765,7 @@ InitializeSmmCpuSemaphores (
   ProcessorCount = gSmmCpuPrivate->SmmCoreEntryContext.NumberOfCpus;\r
   GlobalSemaphoresSize = (sizeof (SMM_CPU_SEMAPHORE_GLOBAL) / sizeof (VOID *)) * SemaphoreSize;\r
   CpuSemaphoresSize    = (sizeof (SMM_CPU_SEMAPHORE_CPU) / sizeof (VOID *)) * ProcessorCount * SemaphoreSize;\r
-  MsrSemahporeSize     = MSR_SPIN_LOCK_INIT_NUM * SemaphoreSize;\r
-  TotalSize = GlobalSemaphoresSize + CpuSemaphoresSize + MsrSemahporeSize;\r
+  TotalSize = GlobalSemaphoresSize + CpuSemaphoresSize;\r
   DEBUG((EFI_D_INFO, "One Semaphore Size    = 0x%x\n", SemaphoreSize));\r
   DEBUG((EFI_D_INFO, "Total Semaphores Size = 0x%x\n", TotalSize));\r
   Pages = EFI_SIZE_TO_PAGES (TotalSize);\r
@@ -1311,12 +1793,6 @@ InitializeSmmCpuSemaphores (
   SemaphoreAddr += ProcessorCount * SemaphoreSize;\r
   mSmmCpuSemaphores.SemaphoreCpu.Present = (BOOLEAN *)SemaphoreAddr;\r
 \r
-  SemaphoreAddr = (UINTN)SemaphoreBlock + GlobalSemaphoresSize + CpuSemaphoresSize;\r
-  mSmmCpuSemaphores.SemaphoreMsr.Msr              = (SPIN_LOCK *)SemaphoreAddr;\r
-  mSmmCpuSemaphores.SemaphoreMsr.AvailableCounter =\r
-        ((UINTN)SemaphoreBlock + Pages * SIZE_4KB - SemaphoreAddr) / SemaphoreSize;\r
-  ASSERT (mSmmCpuSemaphores.SemaphoreMsr.AvailableCounter >= MSR_SPIN_LOCK_INIT_NUM);\r
-\r
   mPFLock                       = mSmmCpuSemaphores.SemaphoreGlobal.PFLock;\r
   mConfigSmmCodeAccessCheckLock = mSmmCpuSemaphores.SemaphoreGlobal.CodeAccessCheckLock;\r
 \r
@@ -1377,14 +1853,16 @@ InitializeMpSyncData (
 /**\r
   Initialize global data for MP synchronization.\r
 \r
-  @param Stacks       Base address of SMI stack buffer for all processors.\r
-  @param StackSize    Stack size for each processor in SMM.\r
+  @param Stacks             Base address of SMI stack buffer for all processors.\r
+  @param StackSize          Stack size for each processor in SMM.\r
+  @param ShadowStackSize    Shadow Stack size for each processor in SMM.\r
 \r
 **/\r
 UINT32\r
 InitializeMpServiceData (\r
   IN VOID        *Stacks,\r
-  IN UINTN       StackSize\r
+  IN UINTN       StackSize,\r
+  IN UINTN       ShadowStackSize\r
   )\r
 {\r
   UINT32                    Cr3;\r
@@ -1436,7 +1914,7 @@ InitializeMpServiceData (
     InstallSmiHandler (\r
       Index,\r
       (UINT32)mCpuHotPlugData.SmBase[Index],\r
-      (VOID*)((UINTN)Stacks + (StackSize * Index)),\r
+      (VOID*)((UINTN)Stacks + (StackSize + ShadowStackSize) * Index),\r
       StackSize,\r
       (UINTN)(GdtTssTables + GdtTableStepSize * Index),\r
       gcSmiGdtr.Limit + 1,\r
@@ -1478,3 +1956,40 @@ RegisterSmmEntry (
   gSmmCpuPrivate->SmmCoreEntry = SmmEntryPoint;\r
   return EFI_SUCCESS;\r
 }\r
+\r
+/**\r
+\r
+  Register the SMM Foundation entry point.\r
+\r
+  @param[in]      Procedure            A pointer to the code stream to be run on the designated target AP\r
+                                       of the system. Type EFI_AP_PROCEDURE is defined below in Volume 2\r
+                                       with the related definitions of\r
+                                       EFI_MP_SERVICES_PROTOCOL.StartupAllAPs.\r
+                                       If caller may pass a value of NULL to deregister any existing\r
+                                       startup procedure.\r
+  @param[in]      ProcedureArguments   Allows the caller to pass a list of parameters to the code that is\r
+                                       run by the AP. It is an optional common mailbox between APs and\r
+                                       the caller to share information\r
+\r
+  @retval EFI_SUCCESS                  The Procedure has been set successfully.\r
+  @retval EFI_INVALID_PARAMETER        The Procedure is NULL but ProcedureArguments not NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+RegisterStartupProcedure (\r
+  IN EFI_AP_PROCEDURE    Procedure,\r
+  IN VOID                *ProcedureArguments OPTIONAL\r
+  )\r
+{\r
+  if (Procedure == NULL && ProcedureArguments != NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+  if (mSmmMpSyncData == NULL) {\r
+    return EFI_NOT_READY;\r
+  }\r
+\r
+  mSmmMpSyncData->StartupProcedure = Procedure;\r
+  mSmmMpSyncData->StartupProcArgs  = ProcedureArguments;\r
+\r
+  return EFI_SUCCESS;\r
+}\r