]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Bus/Ata/AtaAtapiPassThru/AhciMode.c
MdeModulePkg AtaAtapiPassThru: Skip the potential NULL pointer access
[mirror_edk2.git] / MdeModulePkg / Bus / Ata / AtaAtapiPassThru / AhciMode.c
index 3534d9fc2328872ae099d5c85c7c4143cb5d22ea..bda900a1613bf9cc2d7eca78c546f271c7815b62 100644 (file)
@@ -1,15 +1,9 @@
 /** @file\r
   The file for AHCI mode of ATA host controller.\r
 \r
-  Copyright (c) 2010 - 2016, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2010 - 2020, Intel Corporation. All rights reserved.<BR>\r
   (C) Copyright 2015 Hewlett Packard Enterprise Development LP<BR>\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
@@ -225,7 +219,7 @@ AhciWaitMemSet (
 \r
   do {\r
     //\r
-    // Access sytem memory to see if the value is the tested one.\r
+    // Access system memory to see if the value is the tested one.\r
     //\r
     // The system memory pointed by Address will be updated by the\r
     // SATA Host Controller, "volatile" is introduced to prevent\r
@@ -254,32 +248,23 @@ AhciWaitMemSet (
 /**\r
   Check the memory status to the test value.\r
 \r
-  @param[in]       Address           The memory address to test.\r
-  @param[in]       MaskValue         The mask value of memory.\r
-  @param[in]       TestValue         The test value of memory.\r
-  @param[in, out]  Task              Optional. Pointer to the ATA_NONBLOCK_TASK used by\r
-                                     non-blocking mode. If NULL, then just try once.\r
+  @param[in] Address    The memory address to test.\r
+  @param[in] MaskValue  The mask value of memory.\r
+  @param[in] TestValue  The test value of memory.\r
 \r
-  @retval EFI_NOTREADY      The memory is not set.\r
-  @retval EFI_TIMEOUT       The memory setting retry times out.\r
+  @retval EFI_NOT_READY     The memory is not set.\r
   @retval EFI_SUCCESS       The memory is correct set.\r
-\r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
 AhciCheckMemSet (\r
   IN     UINTN                     Address,\r
   IN     UINT32                    MaskValue,\r
-  IN     UINT32                    TestValue,\r
-  IN OUT ATA_NONBLOCK_TASK         *Task\r
+  IN     UINT32                    TestValue\r
   )\r
 {\r
   UINT32     Value;\r
 \r
-  if (Task != NULL) {\r
-    Task->RetryTimes--;\r
-  }\r
-\r
   Value  = *(volatile UINT32 *) Address;\r
   Value &= MaskValue;\r
 \r
@@ -287,49 +272,10 @@ AhciCheckMemSet (
     return EFI_SUCCESS;\r
   }\r
 \r
-  if ((Task != NULL) && !Task->InfiniteWait && (Task->RetryTimes == 0)) {\r
-    return EFI_TIMEOUT;\r
-  } else {\r
-    return EFI_NOT_READY;\r
-  }\r
-}\r
-\r
-/**\r
-  Check if the device is still on port. It also checks if the AHCI controller\r
-  supports the address and data count will be transferred.\r
-\r
-  @param  PciIo            The PCI IO protocol instance.\r
-  @param  Port             The number of port.\r
-\r
-  @retval EFI_SUCCESS      The device is attached to port and the transfer data is\r
-                           supported by AHCI controller.\r
-  @retval EFI_UNSUPPORTED  The transfer address and count is not supported by AHCI\r
-                           controller.\r
-  @retval EFI_NOT_READY    The physical communication between AHCI controller and device\r
-                           is not ready.\r
-\r
-**/\r
-EFI_STATUS\r
-EFIAPI\r
-AhciCheckDeviceStatus (\r
-  IN  EFI_PCI_IO_PROTOCOL    *PciIo,\r
-  IN  UINT8                  Port\r
-  )\r
-{\r
-  UINT32      Data;\r
-  UINT32      Offset;\r
-\r
-  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SSTS;\r
-\r
-  Data   = AhciReadReg (PciIo, Offset) & EFI_AHCI_PORT_SSTS_DET_MASK;\r
-\r
-  if (Data == EFI_AHCI_PORT_SSTS_DET_PCE) {\r
-    return EFI_SUCCESS;\r
-  }\r
-\r
   return EFI_NOT_READY;\r
 }\r
 \r
+\r
 /**\r
 \r
   Clear the port interrupt and error status. It will also clear\r
@@ -371,6 +317,7 @@ AhciClearPortStatus (
   in the Status Register, the Error Register's value is also be dumped.\r
 \r
   @param  PciIo            The PCI IO protocol instance.\r
+  @param  AhciRegisters    The pointer to the EFI_AHCI_REGISTERS.\r
   @param  Port             The number of port.\r
   @param  AtaStatusBlock   A pointer to EFI_ATA_STATUS_BLOCK data structure.\r
 \r
@@ -379,24 +326,42 @@ VOID
 EFIAPI\r
 AhciDumpPortStatus (\r
   IN     EFI_PCI_IO_PROTOCOL        *PciIo,\r
+  IN     EFI_AHCI_REGISTERS         *AhciRegisters,\r
   IN     UINT8                      Port,\r
   IN OUT EFI_ATA_STATUS_BLOCK       *AtaStatusBlock\r
   )\r
 {\r
-  UINT32               Offset;\r
+  UINT               Offset;\r
   UINT32               Data;\r
+  UINTN                FisBaseAddr;\r
+  EFI_STATUS           Status;\r
 \r
   ASSERT (PciIo != NULL);\r
 \r
-  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
-  Data   = AhciReadReg (PciIo, Offset);\r
-\r
   if (AtaStatusBlock != NULL) {\r
     ZeroMem (AtaStatusBlock, sizeof (EFI_ATA_STATUS_BLOCK));\r
 \r
-    AtaStatusBlock->AtaStatus  = (UINT8)Data;\r
-    if ((AtaStatusBlock->AtaStatus & BIT0) != 0) {\r
-      AtaStatusBlock->AtaError = (UINT8)(Data >> 8);\r
+    FisBaseAddr = (UINTN)AhciRegisters->AhciRFis + Port * sizeof (EFI_AHCI_RECEIVED_FIS);\r
+    Offset      = FisBaseAddr + EFI_AHCI_D2H_FIS_OFFSET;\r
+\r
+    Status = AhciCheckMemSet (Offset, EFI_AHCI_FIS_TYPE_MASK, EFI_AHCI_FIS_REGISTER_D2H);\r
+    if (!EFI_ERROR (Status)) {\r
+      //\r
+      // If D2H FIS is received, update StatusBlock with its content.\r
+      //\r
+      CopyMem (AtaStatusBlock, (UINT8 *)Offset, sizeof (EFI_ATA_STATUS_BLOCK));\r
+    } else {\r
+      //\r
+      // If D2H FIS is not received, only update Status & Error field through PxTFD\r
+      // as there is no other way to get the content of the Shadow Register Block.\r
+      //\r
+      Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
+      Data   = AhciReadReg (PciIo, (UINT32)Offset);\r
+\r
+      AtaStatusBlock->AtaStatus  = (UINT8)Data;\r
+      if ((AtaStatusBlock->AtaStatus & BIT0) != 0) {\r
+        AtaStatusBlock->AtaError = (UINT8)(Data >> 8);\r
+      }\r
     }\r
   }\r
 }\r
@@ -427,13 +392,7 @@ AhciEnableFisReceive (
   Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_CMD;\r
   AhciOrReg (PciIo, Offset, EFI_AHCI_PORT_CMD_FRE);\r
 \r
-  return AhciWaitMmioSet (\r
-           PciIo,\r
-           Offset,\r
-           EFI_AHCI_PORT_CMD_FR,\r
-           EFI_AHCI_PORT_CMD_FR,\r
-           Timeout\r
-           );\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -609,7 +568,7 @@ AhciBuildCommand (
 }\r
 \r
 /**\r
-  Buid a command FIS.\r
+  Build a command FIS.\r
 \r
   @param  CmdFis            A pointer to the EFI_AHCI_COMMAND_FIS data structure.\r
   @param  AtaCommandBlock   A pointer to the AhciBuildCommandFis data structure.\r
@@ -649,6 +608,299 @@ AhciBuildCommandFis (
   CmdFis->AhciCFisDevHead     = (UINT8) (AtaCommandBlock->AtaDeviceHead | 0xE0);\r
 }\r
 \r
+/**\r
+  Wait until SATA device reports it is ready for operation.\r
+\r
+  @param[in] PciIo    Pointer to AHCI controller PciIo.\r
+  @param[in] Port     SATA port index on which to reset.\r
+\r
+  @retval EFI_SUCCESS  Device ready for operation.\r
+  @retval EFI_TIMEOUT  Device failed to get ready within required period.\r
+**/\r
+EFI_STATUS\r
+AhciWaitDeviceReady (\r
+  IN EFI_PCI_IO_PROTOCOL  *PciIo,\r
+  IN UINT8                Port\r
+   )\r
+{\r
+  UINT32      PhyDetectDelay;\r
+  UINT32      Data;\r
+  UINT32      Offset;\r
+\r
+  //\r
+  // According to SATA1.0a spec section 5.2, we need to wait for PxTFD.BSY and PxTFD.DRQ\r
+  // and PxTFD.ERR to be zero. The maximum wait time is 16s which is defined at ATA spec.\r
+  //\r
+  PhyDetectDelay = 16 * 1000;\r
+  do {\r
+    Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SERR;\r
+    if (AhciReadReg(PciIo, Offset) != 0) {\r
+      AhciWriteReg (PciIo, Offset, AhciReadReg(PciIo, Offset));\r
+    }\r
+    Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
+\r
+    Data = AhciReadReg (PciIo, Offset) & EFI_AHCI_PORT_TFD_MASK;\r
+    if (Data == 0) {\r
+      break;\r
+    }\r
+\r
+    MicroSecondDelay (1000);\r
+    PhyDetectDelay--;\r
+  } while (PhyDetectDelay > 0);\r
+\r
+  if (PhyDetectDelay == 0) {\r
+    DEBUG ((DEBUG_ERROR, "Port %d Device not ready (TFD=0x%X)\n", Port, Data));\r
+    return EFI_TIMEOUT;\r
+  } else {\r
+    return EFI_SUCCESS;\r
+  }\r
+}\r
+\r
+\r
+/**\r
+  Reset the SATA port. Algorithm follows AHCI spec 1.3.1 section 10.4.2\r
+\r
+  @param[in] PciIo    Pointer to AHCI controller PciIo.\r
+  @param[in] Port     SATA port index on which to reset.\r
+\r
+  @retval EFI_SUCCESS  Port reset.\r
+  @retval Others       Failed to reset the port.\r
+**/\r
+EFI_STATUS\r
+AhciResetPort (\r
+  IN EFI_PCI_IO_PROTOCOL  *PciIo,\r
+  IN UINT8                Port\r
+  )\r
+{\r
+  UINT32      Offset;\r
+  EFI_STATUS  Status;\r
+\r
+  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SCTL;\r
+  AhciOrReg (PciIo, Offset, EFI_AHCI_PORT_SCTL_DET_INIT);\r
+  //\r
+  // SW is required to keep DET set to 0x1 at least for 1 milisecond to ensure that\r
+  // at least one COMRESET signal is sent.\r
+  //\r
+  MicroSecondDelay(1000);\r
+  AhciAndReg (PciIo, Offset, ~(UINT32)EFI_AHCI_PORT_SSTS_DET_MASK);\r
+\r
+  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SSTS;\r
+  Status = AhciWaitMmioSet (PciIo, Offset, EFI_AHCI_PORT_SSTS_DET_MASK, EFI_AHCI_PORT_SSTS_DET_PCE, ATA_ATAPI_TIMEOUT);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
+  return AhciWaitDeviceReady (PciIo, Port);\r
+}\r
+\r
+/**\r
+  Recovers the SATA port from error condition.\r
+  This function implements algorithm described in\r
+  AHCI spec 1.3.1 section 6.2.2\r
+\r
+  @param[in] PciIo    Pointer to AHCI controller PciIo.\r
+  @param[in] Port     SATA port index on which to check.\r
+\r
+  @retval EFI_SUCCESS  Port recovered.\r
+  @retval Others       Failed to recover port.\r
+**/\r
+EFI_STATUS\r
+AhciRecoverPortError (\r
+  IN EFI_PCI_IO_PROTOCOL  *PciIo,\r
+  IN UINT8                Port\r
+  )\r
+{\r
+  UINT32      Offset;\r
+  UINT32      PortInterrupt;\r
+  UINT32      PortTfd;\r
+  EFI_STATUS  Status;\r
+\r
+  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_IS;\r
+  PortInterrupt = AhciReadReg (PciIo, Offset);\r
+  if ((PortInterrupt & EFI_AHCI_PORT_IS_FATAL_ERROR_MASK) == 0) {\r
+    //\r
+    // No fatal error detected. Exit with success as port should still be operational.\r
+    // No need to clear IS as it will be cleared when the next command starts.\r
+    //\r
+    return EFI_SUCCESS;\r
+  }\r
+\r
+  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_CMD;\r
+  AhciAndReg (PciIo, Offset, ~(UINT32)EFI_AHCI_PORT_CMD_ST);\r
+\r
+  Status = AhciWaitMmioSet (PciIo, Offset, EFI_AHCI_PORT_CMD_CR, 0, ATA_ATAPI_TIMEOUT);\r
+  if (EFI_ERROR (Status)) {\r
+    DEBUG ((DEBUG_ERROR, "Ahci port %d is in hung state, aborting recovery\n", Port));\r
+    return Status;\r
+  }\r
+\r
+  //\r
+  // If TFD.BSY or TFD.DRQ is still set it means that drive is hung and software has\r
+  // to reset it before sending any additional commands.\r
+  //\r
+  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
+  PortTfd = AhciReadReg (PciIo, Offset);\r
+  if ((PortTfd & (EFI_AHCI_PORT_TFD_BSY | EFI_AHCI_PORT_TFD_DRQ)) != 0) {\r
+    Status = AhciResetPort (PciIo, Port);\r
+    if (EFI_ERROR (Status)) {\r
+      DEBUG ((DEBUG_ERROR, "Failed to reset the port %d\n", Port));\r
+    }\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Checks if specified FIS has been received.\r
+\r
+  @param[in] PciIo    Pointer to AHCI controller PciIo.\r
+  @param[in] Port     SATA port index on which to check.\r
+  @param[in] FisType  FIS type for which to check.\r
+\r
+  @retval EFI_SUCCESS       FIS received.\r
+  @retval EFI_NOT_READY     FIS not received yet.\r
+  @retval EFI_DEVICE_ERROR  AHCI controller reported an error on port.\r
+**/\r
+EFI_STATUS\r
+AhciCheckFisReceived (\r
+  IN EFI_PCI_IO_PROTOCOL  *PciIo,\r
+  IN UINT8                Port,\r
+  IN SATA_FIS_TYPE        FisType\r
+  )\r
+{\r
+  UINT32      Offset;\r
+  UINT32      PortInterrupt;\r
+  UINT32      PortTfd;\r
+\r
+  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_IS;\r
+  PortInterrupt = AhciReadReg (PciIo, Offset);\r
+  if ((PortInterrupt & EFI_AHCI_PORT_IS_ERROR_MASK) != 0) {\r
+    DEBUG ((DEBUG_ERROR, "AHCI: Error interrupt reported PxIS: %X\n", PortInterrupt));\r
+    return EFI_DEVICE_ERROR;\r
+  }\r
+  //\r
+  // For PIO setup FIS - According to SATA 2.6 spec section 11.7, D2h FIS means an error encountered.\r
+  // But Qemu and Marvel 9230 sata controller may just receive a D2h FIS from device\r
+  // after the transaction is finished successfully.\r
+  // To get better device compatibilities, we further check if the PxTFD's ERR bit is set.\r
+  // By this way, we can know if there is a real error happened.\r
+  //\r
+  if (((FisType == SataFisD2H) && ((PortInterrupt & EFI_AHCI_PORT_IS_DHRS) != 0)) ||\r
+      ((FisType == SataFisPioSetup) && (PortInterrupt & (EFI_AHCI_PORT_IS_PSS | EFI_AHCI_PORT_IS_DHRS)) != 0) ||\r
+      ((FisType == SataFisDmaSetup) && (PortInterrupt & (EFI_AHCI_PORT_IS_DSS | EFI_AHCI_PORT_IS_DHRS)) != 0)) {\r
+    Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
+    PortTfd = AhciReadReg (PciIo, (UINT32) Offset);\r
+    if ((PortTfd & EFI_AHCI_PORT_TFD_ERR) != 0) {\r
+      return EFI_DEVICE_ERROR;\r
+    } else {\r
+      return EFI_SUCCESS;\r
+    }\r
+  }\r
+\r
+  return EFI_NOT_READY;\r
+}\r
+\r
+/**\r
+  Waits until specified FIS has been received.\r
+\r
+  @param[in] PciIo    Pointer to AHCI controller PciIo.\r
+  @param[in] Port     SATA port index on which to check.\r
+  @param[in] Timeout  Time after which function should stop polling.\r
+  @param[in] FisType  FIS type for which to check.\r
+\r
+  @retval EFI_SUCCESS       FIS received.\r
+  @retval EFI_TIMEOUT       FIS failed to arrive within a specified time period.\r
+  @retval EFI_DEVICE_ERROR  AHCI controller reported an error on port.\r
+**/\r
+EFI_STATUS\r
+AhciWaitUntilFisReceived (\r
+  IN EFI_PCI_IO_PROTOCOL  *PciIo,\r
+  IN UINT8                Port,\r
+  IN UINT64               Timeout,\r
+  IN SATA_FIS_TYPE        FisType\r
+  )\r
+{\r
+  EFI_STATUS  Status;\r
+  BOOLEAN     InfiniteWait;\r
+  UINT64      Delay;\r
+\r
+  Delay =  DivU64x32 (Timeout, 1000) + 1;\r
+  if (Timeout == 0) {\r
+    InfiniteWait = TRUE;\r
+  } else {\r
+    InfiniteWait = FALSE;\r
+  }\r
+\r
+  do {\r
+    Status = AhciCheckFisReceived (PciIo, Port, FisType);\r
+    if (Status != EFI_NOT_READY) {\r
+      return Status;\r
+    }\r
+    //\r
+    // Stall for 100 microseconds.\r
+    //\r
+    MicroSecondDelay (100);\r
+    Delay--;\r
+  } while (InfiniteWait || (Delay > 0));\r
+\r
+  return EFI_TIMEOUT;\r
+}\r
+\r
+/**\r
+  Prints contents of the ATA command block into the debug port.\r
+\r
+  @param[in] AtaCommandBlock  AtaCommandBlock to print.\r
+  @param[in] DebugLevel       Debug level on which to print.\r
+**/\r
+VOID\r
+AhciPrintCommandBlock (\r
+  IN EFI_ATA_COMMAND_BLOCK  *AtaCommandBlock,\r
+  IN UINT32                 DebugLevel\r
+  )\r
+{\r
+  DEBUG ((DebugLevel, "ATA COMMAND BLOCK:\n"));\r
+  DEBUG ((DebugLevel, "AtaCommand: %d\n", AtaCommandBlock->AtaCommand));\r
+  DEBUG ((DebugLevel, "AtaFeatures: %X\n", AtaCommandBlock->AtaFeatures));\r
+  DEBUG ((DebugLevel, "AtaSectorNumber: %d\n", AtaCommandBlock->AtaSectorNumber));\r
+  DEBUG ((DebugLevel, "AtaCylinderLow: %X\n", AtaCommandBlock->AtaCylinderHigh));\r
+  DEBUG ((DebugLevel, "AtaCylinderHigh: %X\n", AtaCommandBlock->AtaCylinderHigh));\r
+  DEBUG ((DebugLevel, "AtaDeviceHead: %d\n", AtaCommandBlock->AtaDeviceHead));\r
+  DEBUG ((DebugLevel, "AtaSectorNumberExp: %d\n", AtaCommandBlock->AtaSectorNumberExp));\r
+  DEBUG ((DebugLevel, "AtaCylinderLowExp: %X\n", AtaCommandBlock->AtaCylinderLowExp));\r
+  DEBUG ((DebugLevel, "AtaCylinderHighExp: %X\n", AtaCommandBlock->AtaCylinderHighExp));\r
+  DEBUG ((DebugLevel, "AtaFeaturesExp: %X\n", AtaCommandBlock->AtaFeaturesExp));\r
+  DEBUG ((DebugLevel, "AtaSectorCount: %d\n", AtaCommandBlock->AtaSectorCount));\r
+  DEBUG ((DebugLevel, "AtaSectorCountExp: %d\n", AtaCommandBlock->AtaSectorCountExp));\r
+}\r
+\r
+/**\r
+  Prints contents of the ATA status block into the debug port.\r
+\r
+  @param[in] AtaStatusBlock   AtaStatusBlock to print.\r
+  @param[in] DebugLevel       Debug level on which to print.\r
+**/\r
+VOID\r
+AhciPrintStatusBlock (\r
+  IN EFI_ATA_STATUS_BLOCK  *AtaStatusBlock,\r
+  IN UINT32                DebugLevel\r
+  )\r
+{\r
+  //\r
+  // Skip NULL pointer\r
+  //\r
+  if (AtaStatusBlock == NULL) {\r
+    return;\r
+  }\r
+\r
+  //\r
+  // Only print status and error since we have all of the rest printed as\r
+  // a part of command block print.\r
+  //\r
+  DEBUG ((DebugLevel, "ATA STATUS BLOCK:\n"));\r
+  DEBUG ((DebugLevel, "AtaStatus: %d\n", AtaStatusBlock->AtaStatus));\r
+  DEBUG ((DebugLevel, "AtaError: %d\n", AtaStatusBlock->AtaError));\r
+}\r
+\r
 /**\r
   Start a PIO data transfer on specific port.\r
 \r
@@ -693,26 +945,14 @@ AhciPioTransfer (
   )\r
 {\r
   EFI_STATUS                    Status;\r
-  UINTN                         FisBaseAddr;\r
-  UINTN                         Offset;\r
   EFI_PHYSICAL_ADDRESS          PhyAddr;\r
   VOID                          *Map;\r
   UINTN                         MapLength;\r
   EFI_PCI_IO_PROTOCOL_OPERATION Flag;\r
-  UINT64                        Delay;\r
   EFI_AHCI_COMMAND_FIS          CFis;\r
   EFI_AHCI_COMMAND_LIST         CmdList;\r
-  UINT32                        PortTfd;\r
   UINT32                        PrdCount;\r
-  BOOLEAN                       InfiniteWait;\r
-  BOOLEAN                       PioFisReceived;\r
-  BOOLEAN                       D2hFisReceived;\r
-\r
-  if (Timeout == 0) {\r
-    InfiniteWait = TRUE;\r
-  } else {\r
-    InfiniteWait = FALSE;\r
-  }\r
+  UINT32                        Retry;\r
 \r
   if (Read) {\r
     Flag = EfiPciIoOperationBusMasterWrite;\r
@@ -747,114 +987,58 @@ AhciPioTransfer (
   CmdList.AhciCmdCfl = EFI_AHCI_FIS_REGISTER_H2D_LENGTH / 4;\r
   CmdList.AhciCmdW   = Read ? 0 : 1;\r
 \r
-  AhciBuildCommand (\r
-    PciIo,\r
-    AhciRegisters,\r
-    Port,\r
-    PortMultiplier,\r
-    &CFis,\r
-    &CmdList,\r
-    AtapiCommand,\r
-    AtapiCommandLength,\r
-    0,\r
-    (VOID *)(UINTN)PhyAddr,\r
-    DataCount\r
-    );\r
-\r
-  Status = AhciStartCommand (\r
-             PciIo,\r
-             Port,\r
-             0,\r
-             Timeout\r
-             );\r
-  if (EFI_ERROR (Status)) {\r
-    goto Exit;\r
-  }\r
-\r
-  //\r
-  // Check the status and wait the driver sending data\r
-  //\r
-  FisBaseAddr = (UINTN)AhciRegisters->AhciRFis + Port * sizeof (EFI_AHCI_RECEIVED_FIS);\r
-\r
-  if (Read && (AtapiCommand == 0)) {\r
-    //\r
-    // Wait device sends the PIO setup fis before data transfer\r
-    //\r
-    Status = EFI_TIMEOUT;\r
-    Delay  = DivU64x32 (Timeout, 1000) + 1;\r
-    do {\r
-      PioFisReceived = FALSE;\r
-      D2hFisReceived = FALSE;\r
-      Offset = FisBaseAddr + EFI_AHCI_PIO_FIS_OFFSET;\r
-      Status = AhciCheckMemSet (Offset, EFI_AHCI_FIS_TYPE_MASK, EFI_AHCI_FIS_PIO_SETUP, NULL);\r
-      if (!EFI_ERROR (Status)) {\r
-        PioFisReceived = TRUE;\r
-      }\r
-      //\r
-      // According to SATA 2.6 spec section 11.7, D2h FIS means an error encountered.\r
-      // But Qemu and Marvel 9230 sata controller may just receive a D2h FIS from device\r
-      // after the transaction is finished successfully.\r
-      // To get better device compatibilities, we further check if the PxTFD's ERR bit is set.\r
-      // By this way, we can know if there is a real error happened.\r
-      //\r
-      Offset = FisBaseAddr + EFI_AHCI_D2H_FIS_OFFSET;\r
-      Status = AhciCheckMemSet (Offset, EFI_AHCI_FIS_TYPE_MASK, EFI_AHCI_FIS_REGISTER_D2H, NULL);\r
-      if (!EFI_ERROR (Status)) {\r
-        D2hFisReceived = TRUE;\r
-      }\r
+  for (Retry = 0; Retry < AHCI_COMMAND_RETRIES; Retry++) {\r
+    AhciBuildCommand (\r
+      PciIo,\r
+      AhciRegisters,\r
+      Port,\r
+      PortMultiplier,\r
+      &CFis,\r
+      &CmdList,\r
+      AtapiCommand,\r
+      AtapiCommandLength,\r
+      0,\r
+      (VOID *)(UINTN)PhyAddr,\r
+      DataCount\r
+      );\r
 \r
-      if (PioFisReceived || D2hFisReceived) {\r
-        Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
-        PortTfd = AhciReadReg (PciIo, (UINT32) Offset);\r
-        //\r
-        // PxTFD will be updated if there is a D2H or SetupFIS received. \r
-        //\r
-        if ((PortTfd & EFI_AHCI_PORT_TFD_ERR) != 0) {\r
-          Status = EFI_DEVICE_ERROR;\r
-          break;\r
-        }\r
+    DEBUG ((DEBUG_VERBOSE, "Starting command for PIO transfer:\n"));\r
+    AhciPrintCommandBlock (AtaCommandBlock, DEBUG_VERBOSE);\r
+    Status = AhciStartCommand (\r
+              PciIo,\r
+              Port,\r
+              0,\r
+              Timeout\r
+              );\r
+    if (EFI_ERROR (Status)) {\r
+      break;\r
+    }\r
 \r
+    if (Read && (AtapiCommand == 0)) {\r
+      Status = AhciWaitUntilFisReceived (PciIo, Port, Timeout, SataFisPioSetup);\r
+      if (Status == EFI_SUCCESS) {\r
         PrdCount = *(volatile UINT32 *) (&(AhciRegisters->AhciCmdList[0].AhciCmdPrdbc));\r
         if (PrdCount == DataCount) {\r
           Status = EFI_SUCCESS;\r
-          break;\r
+        } else {\r
+          Status = EFI_DEVICE_ERROR;\r
         }\r
       }\r
-\r
-      //\r
-      // Stall for 100 microseconds.\r
-      //\r
-      MicroSecondDelay(100);\r
-\r
-      Delay--;\r
-      if (Delay == 0) {\r
-        Status = EFI_TIMEOUT;\r
-      }\r
-    } while (InfiniteWait || (Delay > 0));\r
-  } else {\r
-    //\r
-    // Wait for D2H Fis is received\r
-    //\r
-    Offset = FisBaseAddr + EFI_AHCI_D2H_FIS_OFFSET;\r
-    Status = AhciWaitMemSet (\r
-               Offset,\r
-               EFI_AHCI_FIS_TYPE_MASK,\r
-               EFI_AHCI_FIS_REGISTER_D2H,\r
-               Timeout\r
-               );\r
-\r
-    if (EFI_ERROR (Status)) {\r
-      goto Exit;\r
+    } else {\r
+      Status = AhciWaitUntilFisReceived (PciIo, Port, Timeout, SataFisD2H);\r
     }\r
 \r
-    Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
-    PortTfd = AhciReadReg (PciIo, (UINT32) Offset);\r
-    if ((PortTfd & EFI_AHCI_PORT_TFD_ERR) != 0) {\r
-      Status = EFI_DEVICE_ERROR;\r
+    if (Status == EFI_DEVICE_ERROR) {\r
+      DEBUG ((DEBUG_ERROR, "PIO command failed at retry %d\n", Retry));\r
+      Status = AhciRecoverPortError (PciIo, Port);\r
+      if (EFI_ERROR (Status)) {\r
+        break;\r
+      }\r
+    } else {\r
+      break;\r
     }\r
   }\r
 \r
-Exit:\r
   AhciStopCommand (\r
     PciIo,\r
     Port,\r
@@ -872,7 +1056,19 @@ Exit:
     Map\r
     );\r
 \r
-  AhciDumpPortStatus (PciIo, Port, AtaStatusBlock);\r
+  AhciDumpPortStatus (PciIo, AhciRegisters, Port, AtaStatusBlock);\r
+\r
+  if (Status == EFI_DEVICE_ERROR) {\r
+    DEBUG ((DEBUG_ERROR, "Failed to execute command for PIO transfer:\n"));\r
+    //\r
+    // Repeat command block here to make sure it is printed on\r
+    // device error debug level.\r
+    //\r
+    AhciPrintCommandBlock (AtaCommandBlock, DEBUG_ERROR);\r
+    AhciPrintStatusBlock (AtaStatusBlock, DEBUG_ERROR);\r
+  } else {\r
+    AhciPrintStatusBlock (AtaStatusBlock, DEBUG_VERBOSE);\r
+  }\r
 \r
   return Status;\r
 }\r
@@ -921,18 +1117,15 @@ AhciDmaTransfer (
   )\r
 {\r
   EFI_STATUS                    Status;\r
-  UINTN                         Offset;\r
   EFI_PHYSICAL_ADDRESS          PhyAddr;\r
   VOID                          *Map;\r
   UINTN                         MapLength;\r
   EFI_PCI_IO_PROTOCOL_OPERATION Flag;\r
   EFI_AHCI_COMMAND_FIS          CFis;\r
   EFI_AHCI_COMMAND_LIST         CmdList;\r
-  UINTN                         FisBaseAddr;\r
-  UINT32                        PortTfd;\r
-\r
   EFI_PCI_IO_PROTOCOL           *PciIo;\r
   EFI_TPL                       OldTpl;\r
+  UINT32                        Retry;\r
 \r
   Map   = NULL;\r
   PciIo = Instance->PciIo;\r
@@ -942,36 +1135,21 @@ AhciDmaTransfer (
   }\r
 \r
   //\r
-  // Before starting the Blocking BlockIO operation, push to finish all non-blocking\r
-  // BlockIO tasks.\r
-  // Delay 100us to simulate the blocking time out checking.\r
+  // Set Status to suppress incorrect compiler/analyzer warnings\r
   //\r
-  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
-  while ((Task == NULL) && (!IsListEmpty (&Instance->NonBlockingTaskList))) {\r
-    AsyncNonBlockingTransferRoutine (NULL, Instance);\r
-    //\r
-    // Stall for 100us.\r
-    //\r
-    MicroSecondDelay (100);\r
-  }\r
-  gBS->RestoreTPL (OldTpl);\r
+  Status = EFI_SUCCESS;\r
 \r
-  if ((Task == NULL) || ((Task != NULL) && (!Task->IsStart))) {\r
-    //\r
-    // Mark the Task to indicate that it has been started.\r
-    //\r
-    if (Task != NULL) {\r
-      Task->IsStart      = TRUE;\r
-    }\r
+  //\r
+  // DMA buffer allocation. Needs to be done only once for both sync and async\r
+  // DMA transfers irrespective of number of retries.\r
+  //\r
+  if ((Task == NULL) || ((Task != NULL) && (Task->Map == NULL))) {\r
     if (Read) {\r
       Flag = EfiPciIoOperationBusMasterWrite;\r
     } else {\r
       Flag = EfiPciIoOperationBusMasterRead;\r
     }\r
 \r
-    //\r
-    // Construct command list and command table with pci bus address.\r
-    //\r
     MapLength = DataCount;\r
     Status = PciIo->Map (\r
                       PciIo,\r
@@ -985,80 +1163,127 @@ AhciDmaTransfer (
     if (EFI_ERROR (Status) || (DataCount != MapLength)) {\r
       return EFI_BAD_BUFFER_SIZE;\r
     }\r
-\r
     if (Task != NULL) {\r
       Task->Map = Map;\r
     }\r
-    //\r
-    // Package read needed\r
-    //\r
+  }\r
+\r
+  if (Task == NULL || (Task != NULL && !Task->IsStart)) {\r
     AhciBuildCommandFis (&CFis, AtaCommandBlock);\r
 \r
     ZeroMem (&CmdList, sizeof (EFI_AHCI_COMMAND_LIST));\r
 \r
     CmdList.AhciCmdCfl = EFI_AHCI_FIS_REGISTER_H2D_LENGTH / 4;\r
     CmdList.AhciCmdW   = Read ? 0 : 1;\r
-\r
-    AhciBuildCommand (\r
-      PciIo,\r
-      AhciRegisters,\r
-      Port,\r
-      PortMultiplier,\r
-      &CFis,\r
-      &CmdList,\r
-      AtapiCommand,\r
-      AtapiCommandLength,\r
-      0,\r
-      (VOID *)(UINTN)PhyAddr,\r
-      DataCount\r
-      );\r
-\r
-    Status = AhciStartCommand (\r
-               PciIo,\r
-               Port,\r
-               0,\r
-               Timeout\r
-               );\r
-    if (EFI_ERROR (Status)) {\r
-      goto Exit;\r
-    }\r
   }\r
 \r
-  //\r
-  // Wait for command compelte\r
-  //\r
-  FisBaseAddr = (UINTN)AhciRegisters->AhciRFis + Port * sizeof (EFI_AHCI_RECEIVED_FIS);\r
-  Offset      = FisBaseAddr + EFI_AHCI_D2H_FIS_OFFSET;\r
-  if (Task != NULL) {\r
+  if (Task == NULL) {\r
     //\r
-    // For Non-blocking\r
+    // Before starting the Blocking BlockIO operation, push to finish all non-blocking\r
+    // BlockIO tasks.\r
+    // Delay 100us to simulate the blocking time out checking.\r
     //\r
-    Status = AhciCheckMemSet (\r
-               Offset,\r
-               EFI_AHCI_FIS_TYPE_MASK,\r
-               EFI_AHCI_FIS_REGISTER_D2H,\r
-               Task\r
-               );\r
+    OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
+    while (!IsListEmpty (&Instance->NonBlockingTaskList)) {\r
+      AsyncNonBlockingTransferRoutine (NULL, Instance);\r
+      //\r
+      // Stall for 100us.\r
+      //\r
+      MicroSecondDelay (100);\r
+    }\r
+    gBS->RestoreTPL (OldTpl);\r
+    for (Retry = 0; Retry < AHCI_COMMAND_RETRIES; Retry++) {\r
+      AhciBuildCommand (\r
+        PciIo,\r
+        AhciRegisters,\r
+        Port,\r
+        PortMultiplier,\r
+        &CFis,\r
+        &CmdList,\r
+        AtapiCommand,\r
+        AtapiCommandLength,\r
+        0,\r
+        (VOID *)(UINTN)PhyAddr,\r
+        DataCount\r
+        );\r
+\r
+      DEBUG ((DEBUG_VERBOSE, "Starting command for sync DMA transfer:\n"));\r
+      AhciPrintCommandBlock (AtaCommandBlock, DEBUG_VERBOSE);\r
+      Status = AhciStartCommand (\r
+                PciIo,\r
+                Port,\r
+                0,\r
+                Timeout\r
+                );\r
+      if (EFI_ERROR (Status)) {\r
+        break;\r
+      }\r
+      Status = AhciWaitUntilFisReceived (PciIo, Port, Timeout, SataFisD2H);\r
+      if (Status == EFI_DEVICE_ERROR) {\r
+        DEBUG ((DEBUG_ERROR, "DMA command failed at retry: %d\n", Retry));\r
+        Status = AhciRecoverPortError (PciIo, Port);\r
+        if (EFI_ERROR (Status)) {\r
+          break;\r
+        }\r
+      } else {\r
+        break;\r
+      }\r
+    }\r
   } else {\r
-    Status = AhciWaitMemSet (\r
-               Offset,\r
-               EFI_AHCI_FIS_TYPE_MASK,\r
-               EFI_AHCI_FIS_REGISTER_D2H,\r
-               Timeout\r
-               );\r
-  }\r
+    if (!Task->IsStart) {\r
+      AhciBuildCommand (\r
+        PciIo,\r
+        AhciRegisters,\r
+        Port,\r
+        PortMultiplier,\r
+        &CFis,\r
+        &CmdList,\r
+        AtapiCommand,\r
+        AtapiCommandLength,\r
+        0,\r
+        (VOID *)(UINTN)PhyAddr,\r
+        DataCount\r
+        );\r
 \r
-  if (EFI_ERROR (Status)) {\r
-    goto Exit;\r
-  }\r
+      DEBUG ((DEBUG_VERBOSE, "Starting command for async DMA transfer:\n"));\r
+      AhciPrintCommandBlock (AtaCommandBlock, DEBUG_VERBOSE);\r
+      Status = AhciStartCommand (\r
+                PciIo,\r
+                Port,\r
+                0,\r
+                Timeout\r
+                );\r
+      if (!EFI_ERROR (Status)) {\r
+        Task->IsStart = TRUE;\r
+      }\r
+    }\r
+    if (Task->IsStart) {\r
+      Status = AhciCheckFisReceived (PciIo, Port, SataFisD2H);\r
+      if (Status == EFI_DEVICE_ERROR) {\r
+        DEBUG ((DEBUG_ERROR, "DMA command failed at retry: %d\n", Task->RetryTimes));\r
+        Status = AhciRecoverPortError (PciIo, Port);\r
+        //\r
+        // If recovery passed mark the Task as not started and change the status\r
+        // to EFI_NOT_READY. This will make the higher level call this function again\r
+        // and on next call the command will be re-issued due to IsStart being FALSE.\r
+        // This also makes the next condition decrement the RetryTimes.\r
+        //\r
+        if (Status == EFI_SUCCESS) {\r
+          Task->IsStart = FALSE;\r
+          Status = EFI_NOT_READY;\r
+        }\r
+      }\r
 \r
-  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
-  PortTfd = AhciReadReg (PciIo, (UINT32) Offset);\r
-  if ((PortTfd & EFI_AHCI_PORT_TFD_ERR) != 0) {\r
-    Status = EFI_DEVICE_ERROR;\r
+      if (Status == EFI_NOT_READY) {\r
+        if (!Task->InfiniteWait && Task->RetryTimes == 0) {\r
+          Status = EFI_TIMEOUT;\r
+        } else {\r
+          Task->RetryTimes--;\r
+        }\r
+      }\r
+    }\r
   }\r
 \r
-Exit:\r
   //\r
   // For Blocking mode, the command should be stopped, the Fis should be disabled\r
   // and the PciIo should be unmapped.\r
@@ -1091,7 +1316,20 @@ Exit:
     }\r
   }\r
 \r
-  AhciDumpPortStatus (PciIo, Port, AtaStatusBlock);\r
+  AhciDumpPortStatus (PciIo, AhciRegisters, Port, AtaStatusBlock);\r
+\r
+  if (Status == EFI_DEVICE_ERROR) {\r
+    DEBUG ((DEBUG_ERROR, "Failed to execute command for DMA transfer:\n"));\r
+    //\r
+    // Repeat command block here to make sure it is printed on\r
+    // device error debug level.\r
+    //\r
+    AhciPrintCommandBlock (AtaCommandBlock, DEBUG_ERROR);\r
+    AhciPrintStatusBlock (AtaStatusBlock, DEBUG_ERROR);\r
+  } else {\r
+    AhciPrintStatusBlock (AtaStatusBlock, DEBUG_VERBOSE);\r
+  }\r
+\r
   return Status;\r
 }\r
 \r
@@ -1133,11 +1371,9 @@ AhciNonDataTransfer (
   )\r
 {\r
   EFI_STATUS                   Status;\r
-  UINTN                        FisBaseAddr;\r
-  UINTN                        Offset;\r
-  UINT32                       PortTfd;\r
   EFI_AHCI_COMMAND_FIS         CFis;\r
   EFI_AHCI_COMMAND_LIST        CmdList;\r
+  UINT32                       Retry;\r
 \r
   //\r
   // Package read needed\r
@@ -1148,53 +1384,45 @@ AhciNonDataTransfer (
 \r
   CmdList.AhciCmdCfl = EFI_AHCI_FIS_REGISTER_H2D_LENGTH / 4;\r
 \r
-  AhciBuildCommand (\r
-    PciIo,\r
-    AhciRegisters,\r
-    Port,\r
-    PortMultiplier,\r
-    &CFis,\r
-    &CmdList,\r
-    AtapiCommand,\r
-    AtapiCommandLength,\r
-    0,\r
-    NULL,\r
-    0\r
-    );\r
-\r
-  Status = AhciStartCommand (\r
-             PciIo,\r
-             Port,\r
-             0,\r
-             Timeout\r
-             );\r
-  if (EFI_ERROR (Status)) {\r
-    goto Exit;\r
-  }\r
-\r
-  //\r
-  // Wait device sends the Response Fis\r
-  //\r
-  FisBaseAddr = (UINTN)AhciRegisters->AhciRFis + Port * sizeof (EFI_AHCI_RECEIVED_FIS);\r
-  Offset      = FisBaseAddr + EFI_AHCI_D2H_FIS_OFFSET;\r
-  Status      = AhciWaitMemSet (\r
-                  Offset,\r
-                  EFI_AHCI_FIS_TYPE_MASK,\r
-                  EFI_AHCI_FIS_REGISTER_D2H,\r
-                  Timeout\r
-                  );\r
+  for (Retry = 0; Retry < AHCI_COMMAND_RETRIES; Retry++) {\r
+    AhciBuildCommand (\r
+      PciIo,\r
+      AhciRegisters,\r
+      Port,\r
+      PortMultiplier,\r
+      &CFis,\r
+      &CmdList,\r
+      AtapiCommand,\r
+      AtapiCommandLength,\r
+      0,\r
+      NULL,\r
+      0\r
+      );\r
 \r
-  if (EFI_ERROR (Status)) {\r
-    goto Exit;\r
-  }\r
+    DEBUG ((DEBUG_VERBOSE, "Starting command for non data transfer:\n"));\r
+    AhciPrintCommandBlock (AtaCommandBlock, DEBUG_VERBOSE);\r
+    Status = AhciStartCommand (\r
+                PciIo,\r
+                Port,\r
+                0,\r
+                Timeout\r
+                );\r
+    if (EFI_ERROR (Status)) {\r
+      break;\r
+    }\r
 \r
-  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
-  PortTfd = AhciReadReg (PciIo, (UINT32) Offset);\r
-  if ((PortTfd & EFI_AHCI_PORT_TFD_ERR) != 0) {\r
-    Status = EFI_DEVICE_ERROR;\r
+    Status = AhciWaitUntilFisReceived (PciIo, Port, Timeout, SataFisD2H);\r
+    if (Status == EFI_DEVICE_ERROR) {\r
+      DEBUG ((DEBUG_ERROR, "Non data transfer failed at retry %d\n", Retry));\r
+      Status = AhciRecoverPortError (PciIo, Port);\r
+      if (EFI_ERROR (Status)) {\r
+        break;\r
+      }\r
+    } else {\r
+      break;\r
+    }\r
   }\r
 \r
-Exit:\r
   AhciStopCommand (\r
     PciIo,\r
     Port,\r
@@ -1207,7 +1435,19 @@ Exit:
     Timeout\r
     );\r
 \r
-  AhciDumpPortStatus (PciIo, Port, AtaStatusBlock);\r
+  AhciDumpPortStatus (PciIo, AhciRegisters, Port, AtaStatusBlock);\r
+\r
+  if (Status == EFI_DEVICE_ERROR) {\r
+    DEBUG ((DEBUG_ERROR, "Failed to execute command for non data transfer:\n"));\r
+    //\r
+    // Repeat command block here to make sure it is printed on\r
+    // device error debug level.\r
+    //\r
+    AhciPrintCommandBlock (AtaCommandBlock, DEBUG_ERROR);\r
+    AhciPrintStatusBlock (AtaStatusBlock, DEBUG_ERROR);\r
+  } else {\r
+    AhciPrintStatusBlock (AtaStatusBlock, DEBUG_VERBOSE);\r
+  }\r
 \r
   return Status;\r
 }\r
@@ -1336,88 +1576,19 @@ AhciStartCommand (
   }\r
 \r
   Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_CMD;\r
-  AhciOrReg (PciIo, Offset, EFI_AHCI_PORT_CMD_ST | StartCmd);\r
-\r
-  //\r
-  // Setting the command\r
-  //\r
-  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_CI;\r
-  AhciAndReg (PciIo, Offset, 0);\r
-  AhciOrReg (PciIo, Offset, CmdSlotBit);\r
-\r
-  return EFI_SUCCESS;\r
-}\r
-\r
-/**\r
-  Do AHCI port reset.\r
-\r
-  @param  PciIo              The PCI IO protocol instance.\r
-  @param  Port               The number of port.\r
-  @param  Timeout            The timeout value of reset, uses 100ns as a unit.\r
-\r
-  @retval EFI_DEVICE_ERROR   The port reset unsuccessfully\r
-  @retval EFI_TIMEOUT        The reset operation is time out.\r
-  @retval EFI_SUCCESS        The port reset successfully.\r
-\r
-**/\r
-EFI_STATUS\r
-EFIAPI\r
-AhciPortReset (\r
-  IN  EFI_PCI_IO_PROTOCOL       *PciIo,\r
-  IN  UINT8                     Port,\r
-  IN  UINT64                    Timeout\r
-  )\r
-{\r
-  EFI_STATUS      Status;\r
-  UINT32          Offset;\r
-\r
-  AhciClearPortStatus (PciIo, Port);\r
-\r
-  AhciStopCommand (PciIo, Port, Timeout);\r
-\r
-  AhciDisableFisReceive (PciIo, Port, Timeout);\r
-\r
-  AhciEnableFisReceive (PciIo, Port, Timeout);\r
-\r
-  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SCTL;\r
-\r
-  AhciOrReg (PciIo, Offset, EFI_AHCI_PORT_SCTL_DET_INIT);\r
-\r
-  //\r
-  // wait 5 millisecond before de-assert DET\r
-  //\r
-  MicroSecondDelay (5000);\r
-\r
-  AhciAndReg (PciIo, Offset, (UINT32)EFI_AHCI_PORT_SCTL_MASK);\r
-\r
-  //\r
-  // wait 5 millisecond before de-assert DET\r
-  //\r
-  MicroSecondDelay (5000);\r
-\r
-  //\r
-  // Wait for communication to be re-established\r
-  //\r
-  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SSTS;\r
-  Status = AhciWaitMmioSet (\r
-             PciIo,\r
-             Offset,\r
-             EFI_AHCI_PORT_SSTS_DET_MASK,\r
-             EFI_AHCI_PORT_SSTS_DET_PCE,\r
-             Timeout\r
-             );\r
-\r
-  if (EFI_ERROR (Status)) {\r
-    DEBUG ((EFI_D_ERROR, "Port %d COMRESET failed: %r\n", Port, Status));\r
-    return Status;\r
-  }\r
+  AhciOrReg (PciIo, Offset, EFI_AHCI_PORT_CMD_ST | StartCmd);\r
 \r
-  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SERR;\r
-  AhciOrReg (PciIo, Offset, EFI_AHCI_PORT_ERR_CLEAR);\r
+  //\r
+  // Setting the command\r
+  //\r
+  Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_CI;\r
+  AhciAndReg (PciIo, Offset, 0);\r
+  AhciOrReg (PciIo, Offset, CmdSlotBit);\r
 \r
   return EFI_SUCCESS;\r
 }\r
 \r
+\r
 /**\r
   Do AHCI HBA reset.\r
 \r
@@ -1438,17 +1609,13 @@ AhciReset (
 {\r
   UINT64                 Delay;\r
   UINT32                 Value;\r
-  UINT32                 Capability;\r
 \r
   //\r
-  // Collect AHCI controller information\r
-  //\r
-  Capability = AhciReadReg (PciIo, EFI_AHCI_CAPABILITY_OFFSET);\r
-  \r
-  //\r
-  // Enable AE before accessing any AHCI registers if Supports AHCI Mode Only is not set\r
+  // Make sure that GHC.AE bit is set before accessing any AHCI registers.\r
   //\r
-  if ((Capability & EFI_AHCI_CAP_SAM) == 0) {\r
+  Value = AhciReadReg(PciIo, EFI_AHCI_GHC_OFFSET);\r
+\r
+  if ((Value & EFI_AHCI_GHC_ENABLE) == 0) {\r
     AhciOrReg (PciIo, EFI_AHCI_GHC_OFFSET, EFI_AHCI_GHC_ENABLE);\r
   }\r
 \r
@@ -1673,18 +1840,17 @@ AhciAtaSmartSupport (
                    ATA_ATAPI_TIMEOUT,\r
                    NULL\r
                    );\r
-\r
-        if (!EFI_ERROR (Status)) {\r
-          Status = AhciAtaSmartReturnStatusCheck (\r
-                     PciIo,\r
-                     AhciRegisters,\r
-                     (UINT8)Port,\r
-                     (UINT8)PortMultiplier,\r
-                     AtaStatusBlock\r
-                     );\r
-        }\r
       }\r
     }\r
+\r
+    AhciAtaSmartReturnStatusCheck (\r
+      PciIo,\r
+      AhciRegisters,\r
+      (UINT8)Port,\r
+      (UINT8)PortMultiplier,\r
+      AtaStatusBlock\r
+      );\r
+\r
     DEBUG ((EFI_D_INFO, "Enabled S.M.A.R.T feature at port [%d] PortMultiplier [%d]!\n",\r
             Port, PortMultiplier));\r
   }\r
@@ -1817,6 +1983,7 @@ AhciIdentifyPacket (
   @param  PortMultiplier      The port multiplier port number.\r
   @param  Feature             The data to send Feature register.\r
   @param  FeatureSpecificData The specific data for SET FEATURE cmd.\r
+  @param  Timeout             The timeout value of SET FEATURE cmd, uses 100ns as a unit.\r
 \r
   @retval EFI_DEVICE_ERROR    The cmd abort with error occurs.\r
   @retval EFI_TIMEOUT         The operation is time out.\r
@@ -1832,7 +1999,8 @@ AhciDeviceSetFeature (
   IN UINT8                  Port,\r
   IN UINT8                  PortMultiplier,\r
   IN UINT16                 Feature,\r
-  IN UINT32                 FeatureSpecificData\r
+  IN UINT32                 FeatureSpecificData,\r
+  IN UINT64                 Timeout\r
   )\r
 {\r
   EFI_STATUS               Status;\r
@@ -1859,7 +2027,7 @@ AhciDeviceSetFeature (
              0,\r
              &AtaCommandBlock,\r
              &AtaStatusBlock,\r
-             ATA_ATAPI_TIMEOUT,\r
+             Timeout,\r
              NULL\r
              );\r
 \r
@@ -1999,10 +2167,10 @@ AhciCreateTransferDescriptor (
   //\r
   MaxCommandSlotNumber = (UINT8) (((Capability & 0x1F00) >> 8) + 1);\r
   Support64Bit         = (BOOLEAN) (((Capability & BIT31) != 0) ? TRUE : FALSE);\r
-  \r
+\r
   PortImplementBitMap  = AhciReadReg(PciIo, EFI_AHCI_PI_OFFSET);\r
   //\r
-  // Get the highest bit of implemented ports which decides how many bytes are allocated for recived FIS.\r
+  // Get the highest bit of implemented ports which decides how many bytes are allocated for received FIS.\r
   //\r
   MaxPortNumber        = (UINT8)(UINTN)(HighBitSet32(PortImplementBitMap) + 1);\r
   if (MaxPortNumber == 0) {\r
@@ -2057,7 +2225,7 @@ AhciCreateTransferDescriptor (
 \r
   //\r
   // Allocate memory for command list\r
-  // Note that the implemenation is a single task model which only use a command list for all ports.\r
+  // Note that the implementation is a single task model which only use a command list for all ports.\r
   //\r
   Buffer = NULL;\r
   MaxCommandListSize = MaxCommandSlotNumber * sizeof (EFI_AHCI_COMMAND_LIST);\r
@@ -2207,6 +2375,343 @@ Error6:
   return Status;\r
 }\r
 \r
+/**\r
+  Read logs from SATA device.\r
+\r
+  @param  PciIo               The PCI IO protocol instance.\r
+  @param  AhciRegisters       The pointer to the EFI_AHCI_REGISTERS.\r
+  @param  Port                The number of port.\r
+  @param  PortMultiplier      The multiplier of port.\r
+  @param  Buffer              The data buffer to store SATA logs.\r
+  @param  LogNumber           The address of the log.\r
+  @param  PageNumber          The page number of the log.\r
+\r
+  @retval EFI_INVALID_PARAMETER  PciIo, AhciRegisters or Buffer is NULL.\r
+  @retval others                 Return status of AhciPioTransfer().\r
+**/\r
+EFI_STATUS\r
+AhciReadLogExt (\r
+  IN EFI_PCI_IO_PROTOCOL       *PciIo,\r
+  IN EFI_AHCI_REGISTERS        *AhciRegisters,\r
+  IN UINT8                     Port,\r
+  IN UINT8                     PortMultiplier,\r
+  IN OUT UINT8                 *Buffer,\r
+  IN UINT8                     LogNumber,\r
+  IN UINT8                     PageNumber\r
+  )\r
+{\r
+  EFI_ATA_COMMAND_BLOCK        AtaCommandBlock;\r
+  EFI_ATA_STATUS_BLOCK         AtaStatusBlock;\r
+\r
+  if (PciIo == NULL || AhciRegisters == NULL || Buffer == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  ///\r
+  /// Read log from device\r
+  ///\r
+  ZeroMem (&AtaCommandBlock, sizeof (EFI_ATA_COMMAND_BLOCK));\r
+  ZeroMem (&AtaStatusBlock, sizeof (EFI_ATA_STATUS_BLOCK));\r
+  ZeroMem (Buffer, 512);\r
+\r
+  AtaCommandBlock.AtaCommand      = ATA_CMD_READ_LOG_EXT;\r
+  AtaCommandBlock.AtaSectorCount  = 1;\r
+  AtaCommandBlock.AtaSectorNumber = LogNumber;\r
+  AtaCommandBlock.AtaCylinderLow  = PageNumber;\r
+\r
+  return AhciPioTransfer (\r
+           PciIo,\r
+           AhciRegisters,\r
+           Port,\r
+           PortMultiplier,\r
+           NULL,\r
+           0,\r
+           TRUE,\r
+           &AtaCommandBlock,\r
+           &AtaStatusBlock,\r
+           Buffer,\r
+           512,\r
+           ATA_ATAPI_TIMEOUT,\r
+           NULL\r
+           );\r
+}\r
+\r
+/**\r
+  Enable DEVSLP of the disk if supported.\r
+\r
+  @param  PciIo               The PCI IO protocol instance.\r
+  @param  AhciRegisters       The pointer to the EFI_AHCI_REGISTERS.\r
+  @param  Port                The number of port.\r
+  @param  PortMultiplier      The multiplier of port.\r
+  @param  IdentifyData        A pointer to data buffer which is used to contain IDENTIFY data.\r
+\r
+  @retval EFI_SUCCESS         The DEVSLP is enabled per policy successfully.\r
+  @retval EFI_UNSUPPORTED     The DEVSLP isn't supported by the controller/device and policy requires to enable it.\r
+**/\r
+EFI_STATUS\r
+AhciEnableDevSlp (\r
+  IN EFI_PCI_IO_PROTOCOL           *PciIo,\r
+  IN EFI_AHCI_REGISTERS            *AhciRegisters,\r
+  IN UINT8                         Port,\r
+  IN UINT8                         PortMultiplier,\r
+  IN EFI_IDENTIFY_DATA             *IdentifyData\r
+  )\r
+{\r
+  EFI_STATUS               Status;\r
+  UINT32                   Offset;\r
+  UINT32                   Capability2;\r
+  UINT8                    LogData[512];\r
+  DEVSLP_TIMING_VARIABLES  DevSlpTiming;\r
+  UINT32                   PortCmd;\r
+  UINT32                   PortDevSlp;\r
+\r
+  if (mAtaAtapiPolicy->DeviceSleepEnable != 1) {\r
+    return EFI_SUCCESS;\r
+  }\r
+\r
+  //\r
+  // Do not enable DevSlp if DevSlp is not supported.\r
+  //\r
+  Capability2 = AhciReadReg (PciIo, AHCI_CAPABILITY2_OFFSET);\r
+  DEBUG ((DEBUG_INFO, "AHCI CAPABILITY2 = %08x\n", Capability2));\r
+  if ((Capability2 & AHCI_CAP2_SDS) == 0) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  //\r
+  // Do not enable DevSlp if DevSlp is not present\r
+  // Do not enable DevSlp if Hot Plug or Mechanical Presence Switch is supported\r
+  //\r
+  Offset     = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH;\r
+  PortCmd    = AhciReadReg (PciIo, Offset + EFI_AHCI_PORT_CMD);\r
+  PortDevSlp = AhciReadReg (PciIo, Offset + AHCI_PORT_DEVSLP);\r
+  DEBUG ((DEBUG_INFO, "Port CMD/DEVSLP = %08x / %08x\n", PortCmd, PortDevSlp));\r
+  if (((PortDevSlp & AHCI_PORT_DEVSLP_DSP) == 0) ||\r
+      ((PortCmd & (EFI_AHCI_PORT_CMD_HPCP | EFI_AHCI_PORT_CMD_MPSP)) != 0)\r
+     ) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  //\r
+  // Do not enable DevSlp if the device doesn't support DevSlp\r
+  //\r
+  DEBUG ((DEBUG_INFO, "IDENTIFY DEVICE: [77] = %04x, [78] = %04x, [79] = %04x\n",\r
+          IdentifyData->AtaData.reserved_77,\r
+          IdentifyData->AtaData.serial_ata_features_supported, IdentifyData->AtaData.serial_ata_features_enabled));\r
+  if ((IdentifyData->AtaData.serial_ata_features_supported & BIT8) == 0) {\r
+    DEBUG ((DEBUG_INFO, "DevSlp feature is not supported for device at port [%d] PortMultiplier [%d]!\n",\r
+            Port, PortMultiplier));\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  //\r
+  // Enable DevSlp when it is not enabled.\r
+  //\r
+  if ((IdentifyData->AtaData.serial_ata_features_enabled & BIT8) != 0) {\r
+    Status = AhciDeviceSetFeature (\r
+      PciIo, AhciRegisters, Port, 0, ATA_SUB_CMD_ENABLE_SATA_FEATURE, 0x09, ATA_ATAPI_TIMEOUT\r
+    );\r
+    DEBUG ((DEBUG_INFO, "DevSlp set feature for device at port [%d] PortMultiplier [%d] - %r\n",\r
+            Port, PortMultiplier, Status));\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+  }\r
+\r
+  Status = AhciReadLogExt(PciIo, AhciRegisters, Port, PortMultiplier, LogData, 0x30, 0x08);\r
+\r
+  //\r
+  // Clear PxCMD.ST and PxDEVSLP.ADSE before updating PxDEVSLP.DITO and PxDEVSLP.MDAT.\r
+  //\r
+  AhciWriteReg (PciIo, Offset + EFI_AHCI_PORT_CMD, PortCmd & ~EFI_AHCI_PORT_CMD_ST);\r
+  PortDevSlp &= ~AHCI_PORT_DEVSLP_ADSE;\r
+  AhciWriteReg (PciIo, Offset + AHCI_PORT_DEVSLP, PortDevSlp);\r
+\r
+  //\r
+  // Set PxDEVSLP.DETO and PxDEVSLP.MDAT to 0.\r
+  //\r
+  PortDevSlp &= ~AHCI_PORT_DEVSLP_DETO_MASK;\r
+  PortDevSlp &= ~AHCI_PORT_DEVSLP_MDAT_MASK;\r
+  AhciWriteReg (PciIo, Offset + AHCI_PORT_DEVSLP, PortDevSlp);\r
+  DEBUG ((DEBUG_INFO, "Read Log Ext at port [%d] PortMultiplier [%d] - %r\n", Port, PortMultiplier, Status));\r
+  if (EFI_ERROR (Status)) {\r
+    //\r
+    // Assume DEVSLP TIMING VARIABLES is not supported if the Identify Device Data log (30h, 8) fails\r
+    //\r
+    ZeroMem (&DevSlpTiming, sizeof (DevSlpTiming));\r
+  } else {\r
+    CopyMem (&DevSlpTiming, &LogData[48], sizeof (DevSlpTiming));\r
+    DEBUG ((DEBUG_INFO, "DevSlpTiming: Supported(%d), Deto(%d), Madt(%d)\n",\r
+            DevSlpTiming.Supported, DevSlpTiming.Deto, DevSlpTiming.Madt));\r
+  }\r
+\r
+  //\r
+  // Use 20ms as default DETO when DEVSLP TIMING VARIABLES is not supported or the DETO is 0.\r
+  //\r
+  if ((DevSlpTiming.Supported == 0) || (DevSlpTiming.Deto == 0)) {\r
+    DevSlpTiming.Deto = 20;\r
+  }\r
+\r
+  //\r
+  // Use 10ms as default MADT when DEVSLP TIMING VARIABLES is not supported or the MADT is 0.\r
+  //\r
+  if ((DevSlpTiming.Supported == 0) || (DevSlpTiming.Madt == 0)) {\r
+    DevSlpTiming.Madt = 10;\r
+  }\r
+\r
+  PortDevSlp |= DevSlpTiming.Deto << 2;\r
+  PortDevSlp |= DevSlpTiming.Madt << 10;\r
+  AhciOrReg (PciIo, Offset + AHCI_PORT_DEVSLP, PortDevSlp);\r
+\r
+  if (mAtaAtapiPolicy->AggressiveDeviceSleepEnable == 1) {\r
+    if ((Capability2 & AHCI_CAP2_SADM) != 0) {\r
+      PortDevSlp &= ~AHCI_PORT_DEVSLP_DITO_MASK;\r
+      PortDevSlp |= (625 << 15);\r
+      AhciWriteReg (PciIo, Offset + AHCI_PORT_DEVSLP, PortDevSlp);\r
+\r
+      PortDevSlp |= AHCI_PORT_DEVSLP_ADSE;\r
+      AhciWriteReg (PciIo, Offset + AHCI_PORT_DEVSLP, PortDevSlp);\r
+    }\r
+  }\r
+\r
+\r
+  AhciWriteReg (PciIo, Offset + EFI_AHCI_PORT_CMD, PortCmd);\r
+\r
+  DEBUG ((DEBUG_INFO, "Enabled DevSlp feature at port [%d] PortMultiplier [%d], Port CMD/DEVSLP = %08x / %08x\n",\r
+          Port, PortMultiplier, PortCmd, PortDevSlp));\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Spin-up disk if IDD was incomplete or PUIS feature is enabled\r
+\r
+  @param  PciIo               The PCI IO protocol instance.\r
+  @param  AhciRegisters       The pointer to the EFI_AHCI_REGISTERS.\r
+  @param  Port                The number of port.\r
+  @param  PortMultiplier      The multiplier of port.\r
+  @param  IdentifyData        A pointer to data buffer which is used to contain IDENTIFY data.\r
+\r
+**/\r
+EFI_STATUS\r
+AhciSpinUpDisk (\r
+  IN EFI_PCI_IO_PROTOCOL           *PciIo,\r
+  IN EFI_AHCI_REGISTERS            *AhciRegisters,\r
+  IN UINT8                         Port,\r
+  IN UINT8                         PortMultiplier,\r
+  IN OUT EFI_IDENTIFY_DATA         *IdentifyData\r
+  )\r
+{\r
+  EFI_STATUS               Status;\r
+  EFI_ATA_COMMAND_BLOCK    AtaCommandBlock;\r
+  EFI_ATA_STATUS_BLOCK     AtaStatusBlock;\r
+  UINT8                    Buffer[512];\r
+\r
+  if (IdentifyData->AtaData.specific_config == ATA_SPINUP_CFG_REQUIRED_IDD_INCOMPLETE) {\r
+    //\r
+    // Use SET_FEATURE subcommand to spin up the device.\r
+    //\r
+    Status = AhciDeviceSetFeature (\r
+               PciIo, AhciRegisters, Port, PortMultiplier,\r
+               ATA_SUB_CMD_PUIS_SET_DEVICE_SPINUP, 0x00, ATA_SPINUP_TIMEOUT\r
+               );\r
+    DEBUG ((DEBUG_INFO, "CMD_PUIS_SET_DEVICE_SPINUP for device at port [%d] PortMultiplier [%d] - %r!\n",\r
+            Port, PortMultiplier, Status));\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+  } else {\r
+    ASSERT (IdentifyData->AtaData.specific_config == ATA_SPINUP_CFG_NOT_REQUIRED_IDD_INCOMPLETE);\r
+\r
+    //\r
+    // Use READ_SECTORS to spin up the device if SpinUp SET FEATURE subcommand is not supported\r
+    //\r
+    ZeroMem (&AtaCommandBlock, sizeof (EFI_ATA_COMMAND_BLOCK));\r
+    ZeroMem (&AtaStatusBlock, sizeof (EFI_ATA_STATUS_BLOCK));\r
+    //\r
+    // Perform READ SECTORS PIO Data-In command to Read LBA 0\r
+    //\r
+    AtaCommandBlock.AtaCommand      = ATA_CMD_READ_SECTORS;\r
+    AtaCommandBlock.AtaSectorCount  = 0x1;\r
+\r
+    Status = AhciPioTransfer (\r
+               PciIo,\r
+               AhciRegisters,\r
+               Port,\r
+               PortMultiplier,\r
+               NULL,\r
+               0,\r
+               TRUE,\r
+               &AtaCommandBlock,\r
+               &AtaStatusBlock,\r
+               &Buffer,\r
+               sizeof (Buffer),\r
+               ATA_SPINUP_TIMEOUT,\r
+               NULL\r
+               );\r
+    DEBUG ((DEBUG_INFO, "Read LBA 0 for device at port [%d] PortMultiplier [%d] - %r!\n",\r
+            Port, PortMultiplier, Status));\r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
+    }\r
+  }\r
+\r
+  //\r
+  // Read the complete IDENTIFY DEVICE data.\r
+  //\r
+  ZeroMem (IdentifyData, sizeof (*IdentifyData));\r
+  Status = AhciIdentify (PciIo, AhciRegisters, Port, PortMultiplier, IdentifyData);\r
+  if (EFI_ERROR (Status)) {\r
+    DEBUG ((DEBUG_ERROR, "Read IDD failed for device at port [%d] PortMultiplier [%d] - %r!\n",\r
+            Port, PortMultiplier, Status));\r
+    return Status;\r
+  }\r
+\r
+  DEBUG ((DEBUG_INFO, "IDENTIFY DEVICE: [0] = %016x, [2] = %016x, [83] = %016x, [86] = %016x\n",\r
+          IdentifyData->AtaData.config, IdentifyData->AtaData.specific_config,\r
+          IdentifyData->AtaData.command_set_supported_83, IdentifyData->AtaData.command_set_feature_enb_86));\r
+  //\r
+  // Check if IDD is incomplete\r
+  //\r
+  if ((IdentifyData->AtaData.config & BIT2) != 0) {\r
+    return EFI_DEVICE_ERROR;\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Enable/disable/skip PUIS of the disk according to policy.\r
+\r
+  @param  PciIo               The PCI IO protocol instance.\r
+  @param  AhciRegisters       The pointer to the EFI_AHCI_REGISTERS.\r
+  @param  Port                The number of port.\r
+  @param  PortMultiplier      The multiplier of port.\r
+\r
+**/\r
+EFI_STATUS\r
+AhciPuisEnable (\r
+  IN EFI_PCI_IO_PROTOCOL           *PciIo,\r
+  IN EFI_AHCI_REGISTERS            *AhciRegisters,\r
+  IN UINT8                         Port,\r
+  IN UINT8                         PortMultiplier\r
+  )\r
+{\r
+  EFI_STATUS                       Status;\r
+\r
+  Status = EFI_SUCCESS;\r
+  if (mAtaAtapiPolicy->PuisEnable == 0) {\r
+    Status = AhciDeviceSetFeature (PciIo, AhciRegisters, Port, PortMultiplier, ATA_SUB_CMD_DISABLE_PUIS, 0x00, ATA_ATAPI_TIMEOUT);\r
+  } else if (mAtaAtapiPolicy->PuisEnable == 1) {\r
+    Status = AhciDeviceSetFeature (PciIo, AhciRegisters, Port, PortMultiplier, ATA_SUB_CMD_ENABLE_PUIS, 0x00, ATA_ATAPI_TIMEOUT);\r
+  }\r
+  DEBUG ((DEBUG_INFO, "%a PUIS feature at port [%d] PortMultiplier [%d] - %r!\n",\r
+    (mAtaAtapiPolicy->PuisEnable == 0) ? "Disable" : (\r
+    (mAtaAtapiPolicy->PuisEnable == 1) ? "Enable" : "Skip"\r
+      ), Port, PortMultiplier, Status));\r
+  return Status;\r
+}\r
+\r
 /**\r
   Initialize ATA host controller at AHCI mode.\r
 \r
@@ -2239,6 +2744,7 @@ AhciModeInitialization (
   EFI_ATA_COLLECTIVE_MODE          *SupportedModes;\r
   EFI_ATA_TRANSFER_MODE            TransferMode;\r
   UINT32                           PhyDetectDelay;\r
+  UINT32                           Value;\r
 \r
   if (Instance == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
@@ -2257,14 +2763,34 @@ AhciModeInitialization (
   // Collect AHCI controller information\r
   //\r
   Capability = AhciReadReg (PciIo, EFI_AHCI_CAPABILITY_OFFSET);\r
-  \r
+\r
   //\r
-  // Enable AE before accessing any AHCI registers if Supports AHCI Mode Only is not set\r
+  // Make sure that GHC.AE bit is set before accessing any AHCI registers.\r
   //\r
-  if ((Capability & EFI_AHCI_CAP_SAM) == 0) {\r
+  Value = AhciReadReg(PciIo, EFI_AHCI_GHC_OFFSET);\r
+\r
+  if ((Value & EFI_AHCI_GHC_ENABLE) == 0) {\r
     AhciOrReg (PciIo, EFI_AHCI_GHC_OFFSET, EFI_AHCI_GHC_ENABLE);\r
   }\r
-  \r
+\r
+  //\r
+  // Enable 64-bit DMA support in the PCI layer if this controller\r
+  // supports it.\r
+  //\r
+  if ((Capability & EFI_AHCI_CAP_S64A) != 0) {\r
+    Status = PciIo->Attributes (\r
+                      PciIo,\r
+                      EfiPciIoAttributeOperationEnable,\r
+                      EFI_PCI_IO_ATTRIBUTE_DUAL_ADDRESS_CYCLE,\r
+                      NULL\r
+                      );\r
+    if (EFI_ERROR (Status)) {\r
+      DEBUG ((EFI_D_WARN,\r
+        "AhciModeInitialization: failed to enable 64-bit DMA on 64-bit capable controller (%r)\n",\r
+        Status));\r
+    }\r
+  }\r
+\r
   //\r
   // Get the number of command slots per port supported by this HBA.\r
   //\r
@@ -2284,7 +2810,7 @@ AhciModeInitialization (
   }\r
 \r
   for (Port = 0; Port < EFI_AHCI_MAX_PORTS; Port ++) {\r
-    if ((PortImplementBitMap & (BIT0 << Port)) != 0) {\r
+    if ((PortImplementBitMap & (((UINT32)BIT0) << Port)) != 0) {\r
       //\r
       // According to AHCI spec, MaxPortNumber should be equal or greater than the number of implemented ports.\r
       //\r
@@ -2344,20 +2870,9 @@ AhciModeInitialization (
       //\r
       Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_CMD;\r
       AhciOrReg (PciIo, Offset, EFI_AHCI_PORT_CMD_FRE);\r
-      Status = AhciWaitMmioSet (\r
-                 PciIo,\r
-                 Offset,\r
-                 EFI_AHCI_PORT_CMD_FR,\r
-                 EFI_AHCI_PORT_CMD_FR,\r
-                 EFI_AHCI_PORT_CMD_FR_CLEAR_TIMEOUT\r
-                 );\r
-      if (EFI_ERROR (Status)) {\r
-        continue;\r
-      }\r
 \r
       //\r
-      // Wait no longer than 10 ms to wait the Phy to detect the presence of a device.\r
-      // It's the requirment from SATA1.0a spec section 5.2.\r
+      // Wait for the Phy to detect the presence of a device.\r
       //\r
       PhyDetectDelay = EFI_AHCI_BUS_PHY_DETECT_TIMEOUT;\r
       Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SSTS;\r
@@ -2381,29 +2896,8 @@ AhciModeInitialization (
         continue;\r
       }\r
 \r
-      //\r
-      // According to SATA1.0a spec section 5.2, we need to wait for PxTFD.BSY and PxTFD.DRQ\r
-      // and PxTFD.ERR to be zero. The maximum wait time is 16s which is defined at ATA spec.\r
-      //\r
-      PhyDetectDelay = 16 * 1000;\r
-      do {\r
-        Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_SERR;\r
-        if (AhciReadReg(PciIo, Offset) != 0) {\r
-          AhciWriteReg (PciIo, Offset, AhciReadReg(PciIo, Offset));\r
-        }\r
-        Offset = EFI_AHCI_PORT_START + Port * EFI_AHCI_PORT_REG_WIDTH + EFI_AHCI_PORT_TFD;\r
-\r
-        Data = AhciReadReg (PciIo, Offset) & EFI_AHCI_PORT_TFD_MASK;\r
-        if (Data == 0) {\r
-          break;\r
-        }\r
-\r
-        MicroSecondDelay (1000);\r
-        PhyDetectDelay--;\r
-      } while (PhyDetectDelay > 0);\r
-\r
-      if (PhyDetectDelay == 0) {\r
-        DEBUG ((EFI_D_ERROR, "Port %d Device presence detected but phy not ready (TFD=0x%X)\n", Port, Data));\r
+      Status = AhciWaitDeviceReady (PciIo, Port);\r
+      if (EFI_ERROR (Status)) {\r
         continue;\r
       }\r
 \r
@@ -2439,11 +2933,33 @@ AhciModeInitialization (
           continue;\r
         }\r
 \r
+        DEBUG ((\r
+          DEBUG_INFO, "IDENTIFY DEVICE: [0] = %016x, [2] = %016x, [83] = %016x, [86] = %016x\n",\r
+          Buffer.AtaData.config, Buffer.AtaData.specific_config,\r
+          Buffer.AtaData.command_set_supported_83, Buffer.AtaData.command_set_feature_enb_86\r
+          ));\r
+        if ((Buffer.AtaData.config & BIT2) != 0) {\r
+          //\r
+          // SpinUp disk if device reported incomplete IDENTIFY DEVICE.\r
+          //\r
+          Status = AhciSpinUpDisk (\r
+                     PciIo,\r
+                     AhciRegisters,\r
+                     Port,\r
+                     0,\r
+                     &Buffer\r
+                     );\r
+          if (EFI_ERROR (Status)) {\r
+            DEBUG ((DEBUG_ERROR, "Spin up standby device failed - %r\n", Status));\r
+            continue;\r
+          }\r
+        }\r
+\r
         DeviceType = EfiIdeHarddisk;\r
       } else {\r
         continue;\r
       }\r
-      DEBUG ((EFI_D_INFO, "port [%d] port mulitplier [%d] has a [%a]\n",\r
+      DEBUG ((DEBUG_INFO, "port [%d] port multitplier [%d] has a [%a]\n",\r
               Port, 0, DeviceType == EfiIdeCdrom ? "cdrom" : "harddisk"));\r
 \r
       //\r
@@ -2491,7 +3007,7 @@ AhciModeInitialization (
       TransferMode.ModeNumber = (UINT8) (SupportedModes->PioMode.Mode);\r
 \r
       //\r
-      // Set supported DMA mode on this IDE device. Note that UDMA & MDMA cann't\r
+      // Set supported DMA mode on this IDE device. Note that UDMA & MDMA can't\r
       // be set together. Only one DMA mode can be set to a device. If setting\r
       // DMA mode operation fails, we can continue moving on because we only use\r
       // PIO mode at boot time. DMA modes are used by certain kind of OS booting\r
@@ -2504,7 +3020,7 @@ AhciModeInitialization (
         TransferMode.ModeNumber = (UINT8) SupportedModes->MultiWordDmaMode.Mode;\r
       }\r
 \r
-      Status = AhciDeviceSetFeature (PciIo, AhciRegisters, Port, 0, 0x03, (UINT32)(*(UINT8 *)&TransferMode));\r
+      Status = AhciDeviceSetFeature (PciIo, AhciRegisters, Port, 0, 0x03, (UINT32)(*(UINT8 *)&TransferMode), ATA_ATAPI_TIMEOUT);\r
       if (EFI_ERROR (Status)) {\r
         DEBUG ((EFI_D_ERROR, "Set transfer Mode Fail, Status = %r\n", Status));\r
         continue;\r
@@ -2516,6 +3032,29 @@ AhciModeInitialization (
       CreateNewDeviceInfo (Instance, Port, 0xFFFF, DeviceType, &Buffer);\r
       if (DeviceType == EfiIdeHarddisk) {\r
         REPORT_STATUS_CODE (EFI_PROGRESS_CODE, (EFI_PERIPHERAL_FIXED_MEDIA | EFI_P_PC_ENABLE));\r
+        AhciEnableDevSlp (\r
+          PciIo,\r
+          AhciRegisters,\r
+          Port,\r
+          0,\r
+          &Buffer\r
+          );\r
+      }\r
+\r
+      //\r
+      // Enable/disable PUIS according to policy setting if PUIS is capable (Word[83].BIT5 is set).\r
+      //\r
+      if ((Buffer.AtaData.command_set_supported_83 & BIT5) != 0) {\r
+        Status = AhciPuisEnable (\r
+                   PciIo,\r
+                   AhciRegisters,\r
+                   Port,\r
+                   0\r
+                   );\r
+        if (EFI_ERROR (Status)) {\r
+          DEBUG ((DEBUG_ERROR, "PUIS enable/disable failed, Status = %r\n", Status));\r
+          continue;\r
+        }\r
       }\r
     }\r
   }\r