]> git.proxmox.com Git - mirror_edk2.git/blobdiff - UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
UefiCpuPkg/PiSmmCpuDxeSmm: Enable MM MP Protocol
[mirror_edk2.git] / UefiCpuPkg / PiSmmCpuDxeSmm / MpService.c
index 64fb4d6344d5f2aa2ff2fe54ecd1d507996cc443..ef16997547b84e1b36f99ca210c5f4534b73eb21 100644 (file)
@@ -136,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
@@ -347,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
@@ -476,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
@@ -572,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
@@ -604,6 +761,7 @@ APHandler (
   UINT64                            Timer;\r
   UINTN                             BspIndex;\r
   MTRR_SETTINGS                     Mtrrs;\r
+  EFI_STATUS                        ProcedureStatus;\r
 \r
   //\r
   // Timeout BSP\r
@@ -730,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
@@ -906,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
@@ -923,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
@@ -952,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 OUT 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
@@ -995,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
@@ -1020,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,6 +1577,13 @@ SmiRendezvous (
   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
   //\r
@@ -1256,6 +1728,21 @@ Exit:
   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
   Allocate buffer for all semaphores and spin locks.\r
 \r
@@ -1469,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