ArmPsciMpServicesDxe.c

Go to the documentation of this file.

 
#include <PiDxe.h>
 
#include <Library/ArmLib.h>
#include <Library/ArmMmuLib.h>
#include <Library/ArmPlatformLib.h>
#include <Library/ArmSmcLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/CacheMaintenanceLib.h>
#include <Library/CpuExceptionHandlerLib.h>
#include <Library/DebugLib.h>
#include <Library/HobLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiLib.h>
#include <IndustryStandard/ArmStdSmc.h>
#include <Ppi/ArmMpCoreInfo.h>
#include <Protocol/LoadedImage.h>
 
#include "MpServicesInternal.h"
 
#define POLL_INTERVAL_US  50000
 
STATIC CPU_MP_DATA  mCpuMpData;
STATIC BOOLEAN      mNonBlockingModeAllowed;
UINT64              *gApStacksBase;
UINT64              *gProcessorIDs;
CONST UINT64        gApStackSize = AP_STACK_SIZE;
VOID                *gTtbr0;
UINTN               gTcr;
UINTN               gMair;
 
STATIC
BOOLEAN
IsCurrentProcessorBSP (
  VOID
  );
 
STATIC
EFI_STATUS
EFIAPI
DispatchCpu (
  IN UINTN  ProcessorIndex
  )
{
  ARM_SMC_ARGS  Args;
  EFI_STATUS    Status;
 
  Status = EFI_SUCCESS;
 
  mCpuMpData.CpuData[ProcessorIndex].State = CpuStateBusy;
 
  /* Turn the AP on */
  if (sizeof (Args.Arg0) == sizeof (UINT32)) {
    Args.Arg0 = ARM_SMC_ID_PSCI_CPU_ON_AARCH32;
  } else {
    Args.Arg0 = ARM_SMC_ID_PSCI_CPU_ON_AARCH64;
  }
 
  Args.Arg1 = gProcessorIDs[ProcessorIndex];
  Args.Arg2 = (UINTN)ApEntryPoint;
 
  ArmCallSmc (&Args);
 
  if (Args.Arg0 == ARM_SMC_PSCI_RET_ALREADY_ON) {
    Status = EFI_NOT_READY;
  } else if (Args.Arg0 != ARM_SMC_PSCI_RET_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "PSCI_CPU_ON call failed: %d\n", Args.Arg0));
    Status = EFI_DEVICE_ERROR;
  }
 
  return Status;
}
 
STATIC
BOOLEAN
IsProcessorBSP (
  UINTN  ProcessorIndex
  )
{
  EFI_PROCESSOR_INFORMATION  *CpuInfo;
 
  CpuInfo = &mCpuMpData.CpuData[ProcessorIndex].Info;
 
  return (CpuInfo->StatusFlag & PROCESSOR_AS_BSP_BIT) != 0;
}
 
CPU_STATE
GetApState (
  IN  CPU_AP_DATA  *CpuData
  )
{
  return CpuData->State;
}
 
STATIC
VOID
SetApProcedure (
  IN   CPU_AP_DATA       *CpuData,
  IN   EFI_AP_PROCEDURE  Procedure,
  IN   VOID              *ProcedureArgument
  )
{
  ASSERT (CpuData != NULL);
  ASSERT (Procedure != NULL);
 
  CpuData->Parameter = ProcedureArgument;
  CpuData->Procedure = Procedure;
}
 
STATIC
EFI_STATUS
GetNextBlockedNumber (
  OUT UINTN  *NextNumber
  )
{
  UINTN        Index;
  CPU_STATE    State;
  CPU_AP_DATA  *CpuData;
 
  for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
    CpuData = &mCpuMpData.CpuData[Index];
    if (IsProcessorBSP (Index)) {
      // Skip BSP
      continue;
    }
 
    State = CpuData->State;
 
    if (State == CpuStateBlocked) {
      *NextNumber = Index;
      return EFI_SUCCESS;
    }
  }
 
  return EFI_NOT_FOUND;
}
 
STATIC
UINTN
CalculateAndStallInterval (
  IN UINTN  Timeout
  )
{
  UINTN  StallTime;
 
  if ((Timeout < POLL_INTERVAL_US) && (Timeout != 0)) {
    StallTime = Timeout;
  } else {
    StallTime = POLL_INTERVAL_US;
  }
 
  gBS->Stall (StallTime);
 
  return StallTime;
}
 
STATIC
EFI_STATUS
EFIAPI
GetNumberOfProcessors (
  IN  EFI_MP_SERVICES_PROTOCOL  *This,
  OUT UINTN                     *NumberOfProcessors,
  OUT UINTN                     *NumberOfEnabledProcessors
  )
{
  if ((NumberOfProcessors == NULL) || (NumberOfEnabledProcessors == NULL)) {
    return EFI_INVALID_PARAMETER;
  }
 
  if (!IsCurrentProcessorBSP ()) {
    return EFI_DEVICE_ERROR;
  }
 
  *NumberOfProcessors        = mCpuMpData.NumberOfProcessors;
  *NumberOfEnabledProcessors = mCpuMpData.NumberOfEnabledProcessors;
  return EFI_SUCCESS;
}
 
STATIC
EFI_STATUS
EFIAPI
GetProcessorInfo (
  IN  EFI_MP_SERVICES_PROTOCOL   *This,
  IN  UINTN                      ProcessorIndex,
  OUT EFI_PROCESSOR_INFORMATION  *ProcessorInfoBuffer
  )
{
  if (ProcessorInfoBuffer == NULL) {
    return EFI_INVALID_PARAMETER;
  }
 
  if (!IsCurrentProcessorBSP ()) {
    return EFI_DEVICE_ERROR;
  }
 
  ProcessorIndex &= ~CPU_V2_EXTENDED_TOPOLOGY;
 
  if (ProcessorIndex >= mCpuMpData.NumberOfProcessors) {
    return EFI_NOT_FOUND;
  }
 
  CopyMem (
    ProcessorInfoBuffer,
    &mCpuMpData.CpuData[ProcessorIndex].Info,
    sizeof (EFI_PROCESSOR_INFORMATION)
    );
  return EFI_SUCCESS;
}
 
STATIC
EFI_STATUS
EFIAPI
StartupAllAPs (
  IN  EFI_MP_SERVICES_PROTOCOL  *This,
  IN  EFI_AP_PROCEDURE          Procedure,
  IN  BOOLEAN                   SingleThread,
  IN  EFI_EVENT                 WaitEvent               OPTIONAL,
  IN  UINTN                     TimeoutInMicroseconds,
  IN  VOID                      *ProcedureArgument      OPTIONAL,
  OUT UINTN                     **FailedCpuList         OPTIONAL
  )
{
  EFI_STATUS  Status;
 
  if (!IsCurrentProcessorBSP ()) {
    return EFI_DEVICE_ERROR;
  }
 
  if ((mCpuMpData.NumberOfProcessors == 1) || (mCpuMpData.NumberOfEnabledProcessors == 1)) {
    return EFI_NOT_STARTED;
  }
 
  if (Procedure == NULL) {
    return EFI_INVALID_PARAMETER;
  }
 
  if ((WaitEvent != NULL) && !mNonBlockingModeAllowed) {
    return EFI_UNSUPPORTED;
  }
 
  if (FailedCpuList != NULL) {
    mCpuMpData.FailedList = AllocateZeroPool (
                              (mCpuMpData.NumberOfProcessors + 1) *
                              sizeof (UINTN)
                              );
    if (mCpuMpData.FailedList == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }
 
    SetMemN (
      mCpuMpData.FailedList,
      (mCpuMpData.NumberOfProcessors + 1) *
      sizeof (UINTN),
      END_OF_CPU_LIST
      );
    mCpuMpData.FailedListIndex = 0;
    *FailedCpuList             = mCpuMpData.FailedList;
  }
 
  StartupAllAPsPrepareState (SingleThread);
 
  // If any enabled APs are busy (ignoring the BSP), return EFI_NOT_READY
  if (mCpuMpData.StartCount != (mCpuMpData.NumberOfEnabledProcessors - 1)) {
    return EFI_NOT_READY;
  }
 
  if (WaitEvent != NULL) {
    Status = StartupAllAPsWithWaitEvent (
               Procedure,
               ProcedureArgument,
               WaitEvent,
               TimeoutInMicroseconds,
               SingleThread,
               FailedCpuList
               );
 
    if (EFI_ERROR (Status) && (FailedCpuList != NULL)) {
      if (mCpuMpData.FailedListIndex == 0) {
        FreePool (*FailedCpuList);
        *FailedCpuList = NULL;
      }
    }
  } else {
    Status = StartupAllAPsNoWaitEvent (
               Procedure,
               ProcedureArgument,
               TimeoutInMicroseconds,
               SingleThread,
               FailedCpuList
               );
 
    if (FailedCpuList != NULL) {
      if (mCpuMpData.FailedListIndex == 0) {
        FreePool (*FailedCpuList);
        *FailedCpuList = NULL;
      }
    }
  }
 
  return Status;
}
 
STATIC
EFI_STATUS
EFIAPI
StartupThisAP (
  IN  EFI_MP_SERVICES_PROTOCOL  *This,
  IN  EFI_AP_PROCEDURE          Procedure,
  IN  UINTN                     ProcessorNumber,
  IN  EFI_EVENT                 WaitEvent               OPTIONAL,
  IN  UINTN                     TimeoutInMicroseconds,
  IN  VOID                      *ProcedureArgument      OPTIONAL,
  OUT BOOLEAN                   *Finished               OPTIONAL
  )
{
  EFI_STATUS   Status;
  UINTN        Timeout;
  CPU_AP_DATA  *CpuData;
 
  if (!IsCurrentProcessorBSP ()) {
    return EFI_DEVICE_ERROR;
  }
 
  if (Procedure == NULL) {
    return EFI_INVALID_PARAMETER;
  }
 
  if (ProcessorNumber >= mCpuMpData.NumberOfProcessors) {
    return EFI_NOT_FOUND;
  }
 
  CpuData = &mCpuMpData.CpuData[ProcessorNumber];
 
  if (IsProcessorBSP (ProcessorNumber)) {
    return EFI_INVALID_PARAMETER;
  }
 
  if (!IsProcessorEnabled (ProcessorNumber)) {
    return EFI_INVALID_PARAMETER;
  }
 
  if ((GetApState (CpuData) != CpuStateIdle) &&
      (GetApState (CpuData) != CpuStateFinished))
  {
    return EFI_NOT_READY;
  }
 
  if ((WaitEvent != NULL) && !mNonBlockingModeAllowed) {
    return EFI_UNSUPPORTED;
  }
 
  Timeout = TimeoutInMicroseconds;
 
  CpuData->Timeout       = TimeoutInMicroseconds;
  CpuData->TimeTaken     = 0;
  CpuData->TimeoutActive = (BOOLEAN)(TimeoutInMicroseconds != 0);
 
  SetApProcedure (
    CpuData,
    Procedure,
    ProcedureArgument
    );
 
  Status = DispatchCpu (ProcessorNumber);
  if (EFI_ERROR (Status)) {
    CpuData->State = CpuStateIdle;
    return EFI_NOT_READY;
  }
 
  if (WaitEvent != NULL) {
    // Non Blocking
    if (Finished != NULL) {
      CpuData->SingleApFinished = Finished;
      *Finished                 = FALSE;
    }
 
    CpuData->WaitEvent = WaitEvent;
    Status             = gBS->SetTimer (
                                CpuData->CheckThisAPEvent,
                                TimerPeriodic,
                                POLL_INTERVAL_US
                                );
 
    return EFI_SUCCESS;
  }
 
  // Blocking
  while (TRUE) {
    if (GetApState (CpuData) == CpuStateFinished) {
      CpuData->State = CpuStateIdle;
      break;
    }
 
    if ((TimeoutInMicroseconds != 0) && (Timeout == 0)) {
      return EFI_TIMEOUT;
    }
 
    Timeout -= CalculateAndStallInterval (Timeout);
  }
 
  return EFI_SUCCESS;
}
 
STATIC
EFI_STATUS
EFIAPI
SwitchBSP (
  IN EFI_MP_SERVICES_PROTOCOL  *This,
  IN  UINTN                    ProcessorNumber,
  IN  BOOLEAN                  EnableOldBSP
  )
{
  return EFI_UNSUPPORTED;
}
 
STATIC
EFI_STATUS
EFIAPI
EnableDisableAP (
  IN  EFI_MP_SERVICES_PROTOCOL  *This,
  IN  UINTN                     ProcessorNumber,
  IN  BOOLEAN                   EnableAP,
  IN  UINT32                    *HealthFlag OPTIONAL
  )
{
  UINTN        StatusFlag;
  CPU_AP_DATA  *CpuData;
 
  StatusFlag = mCpuMpData.CpuData[ProcessorNumber].Info.StatusFlag;
  CpuData    = &mCpuMpData.CpuData[ProcessorNumber];
 
  if (!IsCurrentProcessorBSP ()) {
    return EFI_DEVICE_ERROR;
  }
 
  if (ProcessorNumber >= mCpuMpData.NumberOfProcessors) {
    return EFI_NOT_FOUND;
  }
 
  if (IsProcessorBSP (ProcessorNumber)) {
    return EFI_INVALID_PARAMETER;
  }
 
  if (GetApState (CpuData) != CpuStateIdle) {
    return EFI_UNSUPPORTED;
  }
 
  if (EnableAP) {
    if (!IsProcessorEnabled (ProcessorNumber)) {
      mCpuMpData.NumberOfEnabledProcessors++;
    }
 
    StatusFlag |= PROCESSOR_ENABLED_BIT;
  } else {
    if (IsProcessorEnabled (ProcessorNumber) && !IsProcessorBSP (ProcessorNumber)) {
      mCpuMpData.NumberOfEnabledProcessors--;
    }
 
    StatusFlag &= ~PROCESSOR_ENABLED_BIT;
  }
 
  if ((HealthFlag != NULL) && !IsProcessorBSP (ProcessorNumber)) {
    StatusFlag &= ~PROCESSOR_HEALTH_STATUS_BIT;
    StatusFlag |= (*HealthFlag & PROCESSOR_HEALTH_STATUS_BIT);
  }
 
  mCpuMpData.CpuData[ProcessorNumber].Info.StatusFlag = StatusFlag;
  return EFI_SUCCESS;
}
 
STATIC
EFI_STATUS
EFIAPI
WhoAmI (
  IN EFI_MP_SERVICES_PROTOCOL  *This,
  OUT UINTN                    *ProcessorNumber
  )
{
  UINTN   Index;
  UINT64  ProcessorId;
 
  if (ProcessorNumber == NULL) {
    return EFI_INVALID_PARAMETER;
  }
 
  ProcessorId = GET_MPIDR_AFFINITY_BITS (ArmReadMpidr ());
  for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
    if (ProcessorId == gProcessorIDs[Index]) {
      *ProcessorNumber = Index;
      break;
    }
  }
 
  return EFI_SUCCESS;
}
 
STATIC EFI_MP_SERVICES_PROTOCOL  mMpServicesProtocol = {
  GetNumberOfProcessors,
  GetProcessorInfo,
  StartupAllAPs,
  StartupThisAP,
  SwitchBSP,
  EnableDisableAP,
  WhoAmI
};
 
STATIC
VOID
AddProcessorToFailedList (
  UINTN      ProcessorIndex,
  CPU_STATE  ApState
  )
{
  UINTN    Index;
  BOOLEAN  Found;
 
  Found = FALSE;
 
  if ((mCpuMpData.FailedList == NULL) ||
      (ApState == CpuStateIdle) ||
      (ApState == CpuStateFinished) ||
      IsProcessorBSP (ProcessorIndex))
  {
    return;
  }
 
  // If we are retrying make sure we don't double count
  for (Index = 0; Index < mCpuMpData.FailedListIndex; Index++) {
    if (mCpuMpData.FailedList[Index] == ProcessorIndex) {
      Found = TRUE;
      break;
    }
  }
 
  /* If the CPU isn't already in the FailedList, add it */
  if (!Found) {
    mCpuMpData.FailedList[mCpuMpData.FailedListIndex++] = ProcessorIndex;
  }
}
 
STATIC
VOID
ProcessStartupAllAPsTimeout (
  VOID
  )
{
  CPU_AP_DATA  *CpuData;
  UINTN        Index;
 
  if (mCpuMpData.FailedList == NULL) {
    return;
  }
 
  for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
    CpuData = &mCpuMpData.CpuData[Index];
    if (IsProcessorBSP (Index)) {
      // Skip BSP
      continue;
    }
 
    if (!IsProcessorEnabled (Index)) {
      // Skip Disabled processors
      continue;
    }
 
    CpuData = &mCpuMpData.CpuData[Index];
    AddProcessorToFailedList (Index, GetApState (CpuData));
  }
}
 
STATIC
VOID
UpdateApStatus (
  IN UINTN  ProcessorIndex
  )
{
  EFI_STATUS   Status;
  CPU_AP_DATA  *CpuData;
  CPU_AP_DATA  *NextCpuData;
  CPU_STATE    State;
  UINTN        NextNumber;
 
  CpuData = &mCpuMpData.CpuData[ProcessorIndex];
 
  if (IsProcessorBSP (ProcessorIndex)) {
    // Skip BSP
    return;
  }
 
  if (!IsProcessorEnabled (ProcessorIndex)) {
    // Skip Disabled processors
    return;
  }
 
  State = GetApState (CpuData);
 
  switch (State) {
    case CpuStateFinished:
      if (mCpuMpData.SingleThread) {
        Status = GetNextBlockedNumber (&NextNumber);
        if (!EFI_ERROR (Status)) {
          NextCpuData = &mCpuMpData.CpuData[NextNumber];
 
          NextCpuData->State = CpuStateReady;
 
          SetApProcedure (
            NextCpuData,
            mCpuMpData.Procedure,
            mCpuMpData.ProcedureArgument
            );
 
          Status = DispatchCpu (NextNumber);
          if (!EFI_ERROR (Status)) {
            mCpuMpData.StartCount++;
          } else {
            AddProcessorToFailedList (NextNumber, NextCpuData->State);
          }
        }
      }
 
      CpuData->State = CpuStateIdle;
      mCpuMpData.FinishCount++;
      break;
 
    default:
      break;
  }
}
 
STATIC
VOID
EFIAPI
CheckAllAPsStatus (
  IN  EFI_EVENT  Event,
  IN  VOID       *Context
  )
{
  EFI_STATUS  Status;
  UINTN       Index;
 
  mCpuMpData.AllTimeTaken += POLL_INTERVAL_US;
 
  for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
    UpdateApStatus (Index);
  }
 
  if (mCpuMpData.AllTimeoutActive && (mCpuMpData.AllTimeTaken > mCpuMpData.AllTimeout)) {
    ProcessStartupAllAPsTimeout ();
 
    // Force terminal exit
    mCpuMpData.FinishCount = mCpuMpData.StartCount;
  }
 
  if (mCpuMpData.FinishCount != mCpuMpData.StartCount) {
    return;
  }
 
  gBS->SetTimer (
         mCpuMpData.CheckAllAPsEvent,
         TimerCancel,
         0
         );
 
  if (mCpuMpData.FailedListIndex == 0) {
    if (mCpuMpData.FailedList != NULL) {
      // Since we don't have the original `FailedCpuList`
      // pointer here to set to NULL, don't free the
      // memory.
    }
  }
 
  Status = gBS->SignalEvent (mCpuMpData.AllWaitEvent);
  ASSERT_EFI_ERROR (Status);
  mCpuMpData.AllWaitEvent = NULL;
}
 
STATIC
VOID
EFIAPI
CheckThisAPStatus (
  IN  EFI_EVENT  Event,
  IN  VOID       *Context
  )
{
  EFI_STATUS   Status;
  CPU_AP_DATA  *CpuData;
  CPU_STATE    State;
 
  CpuData = Context;
 
  CpuData->TimeTaken += POLL_INTERVAL_US;
 
  State = GetApState (CpuData);
 
  if (State == CpuStateFinished) {
    Status = gBS->SetTimer (CpuData->CheckThisAPEvent, TimerCancel, 0);
    ASSERT_EFI_ERROR (Status);
 
    if (CpuData->SingleApFinished != NULL) {
      *(CpuData->SingleApFinished) = TRUE;
    }
 
    if (CpuData->WaitEvent != NULL) {
      Status = gBS->SignalEvent (CpuData->WaitEvent);
      ASSERT_EFI_ERROR (Status);
    }
 
    CpuData->State = CpuStateIdle;
  }
 
  if (CpuData->TimeoutActive && (CpuData->TimeTaken > CpuData->Timeout)) {
    Status = gBS->SetTimer (CpuData->CheckThisAPEvent, TimerCancel, 0);
    if (CpuData->WaitEvent != NULL) {
      Status = gBS->SignalEvent (CpuData->WaitEvent);
      ASSERT_EFI_ERROR (Status);
      CpuData->WaitEvent = NULL;
    }
  }
}
 
STATIC
EFI_STATUS
FillInProcessorInformation (
  IN BOOLEAN  BSP,
  IN UINTN    Mpidr,
  IN UINTN    ProcessorIndex
  )
{
  EFI_PROCESSOR_INFORMATION  *CpuInfo;
 
  CpuInfo = &mCpuMpData.CpuData[ProcessorIndex].Info;
 
  CpuInfo->ProcessorId = GET_MPIDR_AFFINITY_BITS (Mpidr);
  CpuInfo->StatusFlag  = PROCESSOR_ENABLED_BIT | PROCESSOR_HEALTH_STATUS_BIT;
 
  if (BSP) {
    CpuInfo->StatusFlag |= PROCESSOR_AS_BSP_BIT;
  }
 
  if ((Mpidr & MPIDR_MT_BIT) > 0) {
    CpuInfo->Location.Package = GET_MPIDR_AFF2 (Mpidr);
    CpuInfo->Location.Core    = GET_MPIDR_AFF1 (Mpidr);
    CpuInfo->Location.Thread  = GET_MPIDR_AFF0 (Mpidr);
 
    CpuInfo->ExtendedInformation.Location2.Package = GET_MPIDR_AFF3 (Mpidr);
    CpuInfo->ExtendedInformation.Location2.Die     = GET_MPIDR_AFF2 (Mpidr);
    CpuInfo->ExtendedInformation.Location2.Core    = GET_MPIDR_AFF1 (Mpidr);
    CpuInfo->ExtendedInformation.Location2.Thread  = GET_MPIDR_AFF0 (Mpidr);
  } else {
    CpuInfo->Location.Package = GET_MPIDR_AFF1 (Mpidr);
    CpuInfo->Location.Core    = GET_MPIDR_AFF0 (Mpidr);
    CpuInfo->Location.Thread  = 0;
 
    CpuInfo->ExtendedInformation.Location2.Package = GET_MPIDR_AFF2 (Mpidr);
    CpuInfo->ExtendedInformation.Location2.Die     = GET_MPIDR_AFF1 (Mpidr);
    CpuInfo->ExtendedInformation.Location2.Core    = GET_MPIDR_AFF0 (Mpidr);
    CpuInfo->ExtendedInformation.Location2.Thread  = 0;
  }
 
  mCpuMpData.CpuData[ProcessorIndex].State = BSP ? CpuStateBusy : CpuStateIdle;
 
  mCpuMpData.CpuData[ProcessorIndex].Procedure = NULL;
  mCpuMpData.CpuData[ProcessorIndex].Parameter = NULL;
 
  return EFI_SUCCESS;
}
 
STATIC
EFI_STATUS
MpServicesInitialize (
  IN   UINTN                NumberOfProcessors,
  IN   CONST ARM_CORE_INFO  *CoreInfo
  )
{
  EFI_STATUS  Status;
  UINTN       Index;
  EFI_EVENT   ReadyToBootEvent;
  BOOLEAN     IsBsp;
 
  //
  // Clear the data structure area first.
  //
  ZeroMem (&mCpuMpData, sizeof (CPU_MP_DATA));
  //
  // First BSP fills and inits all known values, including its own records.
  //
  mCpuMpData.NumberOfProcessors        = NumberOfProcessors;
  mCpuMpData.NumberOfEnabledProcessors = NumberOfProcessors;
 
  mCpuMpData.CpuData = AllocateZeroPool (
                         mCpuMpData.NumberOfProcessors * sizeof (CPU_AP_DATA)
                         );
 
  if (mCpuMpData.CpuData == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
 
  /* Allocate one extra for the sentinel entry at the end */
  gProcessorIDs = AllocateZeroPool ((mCpuMpData.NumberOfProcessors + 1) * sizeof (UINT64));
  ASSERT (gProcessorIDs != NULL);
 
  Status = gBS->CreateEvent (
                  EVT_TIMER | EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  CheckAllAPsStatus,
                  NULL,
                  &mCpuMpData.CheckAllAPsEvent
                  );
  ASSERT_EFI_ERROR (Status);
 
  gApStacksBase = AllocatePages (
                    EFI_SIZE_TO_PAGES (
                      mCpuMpData.NumberOfProcessors *
                      gApStackSize
                      )
                    );
  ASSERT (gApStacksBase != NULL);
 
  for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
    if (GET_MPIDR_AFFINITY_BITS (ArmReadMpidr ()) == CoreInfo[Index].Mpidr) {
      IsBsp = TRUE;
    } else {
      IsBsp = FALSE;
    }
 
    FillInProcessorInformation (IsBsp, CoreInfo[Index].Mpidr, Index);
 
    gProcessorIDs[Index] = mCpuMpData.CpuData[Index].Info.ProcessorId;
 
    Status = gBS->CreateEvent (
                    EVT_TIMER | EVT_NOTIFY_SIGNAL,
                    TPL_CALLBACK,
                    CheckThisAPStatus,
                    (VOID *)&mCpuMpData.CpuData[Index],
                    &mCpuMpData.CpuData[Index].CheckThisAPEvent
                    );
    ASSERT_EFI_ERROR (Status);
  }
 
  gProcessorIDs[Index] = MAX_UINT64;
 
  gTcr   = ArmGetTCR ();
  gMair  = ArmGetMAIR ();
  gTtbr0 = ArmGetTTBR0BaseAddress ();
 
  //
  // The global pointer variables as well as the gProcessorIDs array contents
  // are accessed by the other cores so we must clean them to the PoC
  //
  WriteBackDataCacheRange (&gProcessorIDs, sizeof (UINT64 *));
  WriteBackDataCacheRange (&gApStacksBase, sizeof (UINT64 *));
 
  WriteBackDataCacheRange (
    gProcessorIDs,
    (mCpuMpData.NumberOfProcessors + 1) * sizeof (UINT64)
    );
 
  mNonBlockingModeAllowed = TRUE;
 
  Status = EfiCreateEventReadyToBootEx (
             TPL_CALLBACK,
             ReadyToBootSignaled,
             NULL,
             &ReadyToBootEvent
             );
  ASSERT_EFI_ERROR (Status);
 
  return EFI_SUCCESS;
}
 
STATIC
VOID
EFIAPI
ReadyToBootSignaled (
  IN  EFI_EVENT  Event,
  IN  VOID       *Context
  )
{
  mNonBlockingModeAllowed = FALSE;
}
 
EFI_STATUS
EFIAPI
ArmPsciMpServicesDxeInitialize (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS                 Status;
  EFI_HANDLE                 Handle;
  UINTN                      MaxCpus;
  EFI_LOADED_IMAGE_PROTOCOL  *Image;
  EFI_HOB_GENERIC_HEADER     *Hob;
  VOID                       *HobData;
  UINTN                      HobDataSize;
  CONST ARM_CORE_INFO        *CoreInfo;
 
  MaxCpus = 1;
 
  Status = gBS->HandleProtocol (
                  ImageHandle,
                  &gEfiLoadedImageProtocolGuid,
                  (VOID **)&Image
                  );
  ASSERT_EFI_ERROR (Status);
 
  //
  // Parts of the code in this driver may be executed by other cores running
  // with the MMU off so we need to ensure that everything is clean to the
  // point of coherency (PoC)
  //
  WriteBackDataCacheRange (Image->ImageBase, Image->ImageSize);
 
  Hob = GetFirstGuidHob (&gArmMpCoreInfoGuid);
  if (Hob != NULL) {
    HobData     = GET_GUID_HOB_DATA (Hob);
    HobDataSize = GET_GUID_HOB_DATA_SIZE (Hob);
    CoreInfo    = (ARM_CORE_INFO *)HobData;
    MaxCpus     = HobDataSize / sizeof (ARM_CORE_INFO);
  }
 
  if (MaxCpus == 1) {
    DEBUG ((DEBUG_WARN, "Trying to use EFI_MP_SERVICES_PROTOCOL on a UP system"));
    // We are not MP so nothing to do
    return EFI_NOT_FOUND;
  }
 
  Status = MpServicesInitialize (MaxCpus, CoreInfo);
  if (Status != EFI_SUCCESS) {
    ASSERT_EFI_ERROR (Status);
    return Status;
  }
 
  //
  // Now install the MP services protocol.
  //
  Handle = NULL;
  Status = gBS->InstallMultipleProtocolInterfaces (
                  &Handle,
                  &gEfiMpServiceProtocolGuid,
                  &mMpServicesProtocol,
                  NULL
                  );
  ASSERT_EFI_ERROR (Status);
 
  return Status;
}
 
STATIC
VOID
EFIAPI
ApExceptionHandler (
  IN CONST EFI_EXCEPTION_TYPE  InterruptType,
  IN CONST EFI_SYSTEM_CONTEXT  SystemContext
  )
{
  ARM_SMC_ARGS  Args;
  UINT64        Mpidr;
  UINTN         Index;
  UINTN         ProcessorIndex;
 
  Mpidr = GET_MPIDR_AFFINITY_BITS (ArmReadMpidr ());
 
  Index          = 0;
  ProcessorIndex = MAX_UINT64;
 
  do {
    if (gProcessorIDs[Index] == Mpidr) {
      ProcessorIndex = Index;
      break;
    }
 
    Index++;
  } while (gProcessorIDs[Index] != MAX_UINT64);
 
  if (ProcessorIndex != MAX_UINT64) {
    mCpuMpData.CpuData[ProcessorIndex].State = CpuStateFinished;
    ArmDataMemoryBarrier ();
  }
 
  Args.Arg0 = ARM_SMC_ID_PSCI_CPU_OFF;
  ArmCallSmc (&Args);
 
  /* Should never be reached */
  ASSERT (FALSE);
  CpuDeadLoop ();
}
 
VOID
ApProcedure (
  VOID
  )
{
  ARM_SMC_ARGS      Args;
  EFI_AP_PROCEDURE  UserApProcedure;
  VOID              *UserApParameter;
  UINTN             ProcessorIndex;
 
  ProcessorIndex = 0;
 
  WhoAmI (&mMpServicesProtocol, &ProcessorIndex);
 
  /* Fetch the user-supplied procedure and parameter to execute */
  UserApProcedure = mCpuMpData.CpuData[ProcessorIndex].Procedure;
  UserApParameter = mCpuMpData.CpuData[ProcessorIndex].Parameter;
 
  InitializeCpuExceptionHandlers (NULL);
  RegisterCpuInterruptHandler (EXCEPT_AARCH64_SYNCHRONOUS_EXCEPTIONS, ApExceptionHandler);
  RegisterCpuInterruptHandler (EXCEPT_AARCH64_IRQ, ApExceptionHandler);
  RegisterCpuInterruptHandler (EXCEPT_AARCH64_FIQ, ApExceptionHandler);
  RegisterCpuInterruptHandler (EXCEPT_AARCH64_SERROR, ApExceptionHandler);
 
  UserApProcedure (UserApParameter);
 
  mCpuMpData.CpuData[ProcessorIndex].State = CpuStateFinished;
 
  ArmDataMemoryBarrier ();
 
  /* Since we're finished with this AP, turn it off */
  Args.Arg0 = ARM_SMC_ID_PSCI_CPU_OFF;
  ArmCallSmc (&Args);
 
  /* Should never be reached */
  ASSERT (FALSE);
  CpuDeadLoop ();
}
 
STATIC
BOOLEAN
IsCurrentProcessorBSP (
  VOID
  )
{
  EFI_STATUS  Status;
  UINTN       ProcessorIndex;
 
  Status = WhoAmI (&mMpServicesProtocol, &ProcessorIndex);
  if (EFI_ERROR (Status)) {
    ASSERT_EFI_ERROR (Status);
    return FALSE;
  }
 
  return IsProcessorBSP (ProcessorIndex);
}
 
STATIC
BOOLEAN
IsProcessorEnabled (
  UINTN  ProcessorIndex
  )
{
  EFI_PROCESSOR_INFORMATION  *CpuInfo;
 
  CpuInfo = &mCpuMpData.CpuData[ProcessorIndex].Info;
 
  return (CpuInfo->StatusFlag & PROCESSOR_ENABLED_BIT) != 0;
}
 
STATIC
VOID
StartupAllAPsPrepareState (
  IN BOOLEAN  SingleThread
  )
{
  UINTN        Index;
  CPU_STATE    APInitialState;
  CPU_AP_DATA  *CpuData;
 
  mCpuMpData.FinishCount  = 0;
  mCpuMpData.StartCount   = 0;
  mCpuMpData.SingleThread = SingleThread;
 
  APInitialState = CpuStateReady;
 
  for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
    CpuData = &mCpuMpData.CpuData[Index];
 
    //
    // Get APs prepared, and put failing APs into FailedCpuList.
    // If "SingleThread", only 1 AP will put into ready state, other AP will be
    // put into ready state 1 by 1, until the previous 1 finished its task.
    // If not "SingleThread", all APs are put into ready state from the
    // beginning
    //
 
    if (IsProcessorBSP (Index)) {
      // Skip BSP
      continue;
    }
 
    if (!IsProcessorEnabled (Index)) {
      // Skip Disabled processors
      if (mCpuMpData.FailedList != NULL) {
        mCpuMpData.FailedList[mCpuMpData.FailedListIndex++] = Index;
      }
 
      continue;
    }
 
    // If any APs finished after timing out, reset state to Idle
    if (GetApState (CpuData) == CpuStateFinished) {
      CpuData->State = CpuStateIdle;
    }
 
    if (GetApState (CpuData) != CpuStateIdle) {
      // Skip busy processors
      if (mCpuMpData.FailedList != NULL) {
        mCpuMpData.FailedList[mCpuMpData.FailedListIndex++] = Index;
      }
    }
 
    CpuData->State = APInitialState;
 
    mCpuMpData.StartCount++;
    if (SingleThread) {
      APInitialState = CpuStateBlocked;
    }
  }
}
 
STATIC
EFI_STATUS
StartupAllAPsWithWaitEvent (
  IN EFI_AP_PROCEDURE  Procedure,
  IN VOID              *ProcedureArgument,
  IN EFI_EVENT         WaitEvent,
  IN UINTN             TimeoutInMicroseconds,
  IN BOOLEAN           SingleThread,
  IN UINTN             **FailedCpuList
  )
{
  EFI_STATUS   Status;
  UINTN        Index;
  CPU_AP_DATA  *CpuData;
 
  for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
    CpuData = &mCpuMpData.CpuData[Index];
    if (IsProcessorBSP (Index)) {
      // Skip BSP
      continue;
    }
 
    if (!IsProcessorEnabled (Index)) {
      // Skip Disabled processors
      continue;
    }
 
    if (GetApState (CpuData) == CpuStateReady) {
      SetApProcedure (CpuData, Procedure, ProcedureArgument);
      if ((mCpuMpData.StartCount == 0) || !SingleThread) {
        Status = DispatchCpu (Index);
        if (EFI_ERROR (Status)) {
          AddProcessorToFailedList (Index, CpuData->State);
          break;
        }
      }
    }
  }
 
  if (EFI_ERROR (Status)) {
    return EFI_NOT_READY;
  }
 
  //
  // Save data into private data structure, and create timer to poll AP state
  // before exiting
  //
  mCpuMpData.Procedure         = Procedure;
  mCpuMpData.ProcedureArgument = ProcedureArgument;
  mCpuMpData.AllWaitEvent      = WaitEvent;
  mCpuMpData.AllTimeout        = TimeoutInMicroseconds;
  mCpuMpData.AllTimeTaken      = 0;
  mCpuMpData.AllTimeoutActive  = (BOOLEAN)(TimeoutInMicroseconds != 0);
  Status                       = gBS->SetTimer (
                                        mCpuMpData.CheckAllAPsEvent,
                                        TimerPeriodic,
                                        POLL_INTERVAL_US
                                        );
 
  return Status;
}
 
STATIC
EFI_STATUS
StartupAllAPsNoWaitEvent (
  IN EFI_AP_PROCEDURE  Procedure,
  IN VOID              *ProcedureArgument,
  IN UINTN             TimeoutInMicroseconds,
  IN BOOLEAN           SingleThread,
  IN UINTN             **FailedCpuList
  )
{
  EFI_STATUS   Status;
  UINTN        Index;
  UINTN        NextIndex;
  UINTN        Timeout;
  CPU_AP_DATA  *CpuData;
  BOOLEAN      DispatchError;
 
  Timeout       = TimeoutInMicroseconds;
  DispatchError = FALSE;
 
  while (TRUE) {
    for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
      CpuData = &mCpuMpData.CpuData[Index];
      if (IsProcessorBSP (Index)) {
        // Skip BSP
        continue;
      }
 
      if (!IsProcessorEnabled (Index)) {
        // Skip Disabled processors
        continue;
      }
 
      switch (GetApState (CpuData)) {
        case CpuStateReady:
          SetApProcedure (CpuData, Procedure, ProcedureArgument);
          Status = DispatchCpu (Index);
          if (EFI_ERROR (Status)) {
            AddProcessorToFailedList (Index, CpuData->State);
            CpuData->State = CpuStateIdle;
            mCpuMpData.StartCount--;
            DispatchError = TRUE;
 
            if (SingleThread) {
              // Dispatch the next available AP
              Status = GetNextBlockedNumber (&NextIndex);
              if (!EFI_ERROR (Status)) {
                mCpuMpData.CpuData[NextIndex].State = CpuStateReady;
              }
            }
          }
 
          break;
 
        case CpuStateFinished:
          mCpuMpData.FinishCount++;
          if (SingleThread) {
            Status = GetNextBlockedNumber (&NextIndex);
            if (!EFI_ERROR (Status)) {
              mCpuMpData.CpuData[NextIndex].State = CpuStateReady;
            }
          }
 
          CpuData->State = CpuStateIdle;
          break;
 
        default:
          break;
      }
    }
 
    if (mCpuMpData.FinishCount == mCpuMpData.StartCount) {
      Status = EFI_SUCCESS;
      break;
    }
 
    if ((TimeoutInMicroseconds != 0) && (Timeout == 0)) {
      Status = EFI_TIMEOUT;
      break;
    }
 
    Timeout -= CalculateAndStallInterval (Timeout);
  }
 
  if (Status == EFI_TIMEOUT) {
    // Add any remaining CPUs to the FailedCpuList
    if (FailedCpuList != NULL) {
      for (Index = 0; Index < mCpuMpData.NumberOfProcessors; Index++) {
        AddProcessorToFailedList (Index, mCpuMpData.CpuData[Index].State);
      }
    }
  }
 
  if (DispatchError) {
    Status = EFI_NOT_READY;
  }
 
  return Status;
}