text: /** @file
  
    Copyright (c) 2014 - 2021, Intel Corporation. All rights reserved.<BR>
    SPDX-License-Identifier: BSD-2-Clause-Patent
  
  **/
  
  #include "UfsBlockIoPei.h"
  
  /**
    Wait for the value of the specified system memory set to the test value.
  
    @param  Address           The system memory address to test.
    @param  MaskValue         The mask value of memory.
    @param  TestValue         The test value of memory.
    @param  Timeout           The time out value for wait memory set, uses 100ns as a unit.
  
    @retval EFI_TIMEOUT       The system memory setting is time out.
    @retval EFI_SUCCESS       The system memory is correct set.
  
  **/
  EFI_STATUS
  EFIAPI
  UfsWaitMemSet (
    IN  UINTN   Address,
    IN  UINT32  MaskValue,
    IN  UINT32  TestValue,
    IN  UINT64  Timeout
    )
  {
    UINT32   Value;
    UINT64   Delay;
    BOOLEAN  InfiniteWait;
  
    if (Timeout == 0) {
      InfiniteWait = TRUE;
    } else {
      InfiniteWait = FALSE;
    }
  
    Delay = DivU64x32 (Timeout, 10) + 1;
  
    do {
      //
      // Access PCI MMIO space to see if the value is the tested one.
      //
      Value = MmioRead32 (Address) & MaskValue;
  
      if (Value == TestValue) {
        return EFI_SUCCESS;
      }
  
      //
      // Stall for 1 microseconds.
      //
      MicroSecondDelay (1);
  
      Delay--;
    } while (InfiniteWait || (Delay > 0));
  
    return EFI_TIMEOUT;
  }
  
  /**
    Dump UIC command execution result for debugging.
  
    @param[in]   UicOpcode  The executed UIC opcode.
    @param[in]   Result     The result to be parsed.
  
  **/
  VOID
  DumpUicCmdExecResult (
    IN  UINT8  UicOpcode,
    IN  UINT8  Result
    )
  {
    if (UicOpcode <= UfsUicDmePeerSet) {
      switch (Result) {
        case 0x00:
          break;
        case 0x01:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - INVALID_MIB_ATTRIBUTE\n"));
          break;
        case 0x02:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - INVALID_MIB_ATTRIBUTE_VALUE\n"));
          break;
        case 0x03:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - READ_ONLY_MIB_ATTRIBUTE\n"));
          break;
        case 0x04:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - WRITE_ONLY_MIB_ATTRIBUTE\n"));
          break;
        case 0x05:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - BAD_INDEX\n"));
          break;
        case 0x06:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - LOCKED_MIB_ATTRIBUTE\n"));
          break;
        case 0x07:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - BAD_TEST_FEATURE_INDEX\n"));
          break;
        case 0x08:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - PEER_COMMUNICATION_FAILURE\n"));
          break;
        case 0x09:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - BUSY\n"));
          break;
        case 0x0A:
          DEBUG ((DEBUG_VERBOSE, "UIC configuration command fails - DME_FAILURE\n"));
          break;
        default:
          ASSERT (FALSE);
          break;
      }
    } else {
      switch (Result) {
        case 0x00:
          break;
        case 0x01:
          DEBUG ((DEBUG_VERBOSE, "UIC control command fails - FAILURE\n"));
          break;
        default:
          ASSERT (FALSE);
          break;
      }
    }
  }
  
  /**
    Dump QUERY RESPONSE UPIU result for debugging.
  
    @param[in]   Result  The result to be parsed.
  
  **/
  VOID
  DumpQueryResponseResult (
    IN  UINT8  Result
    )
  {
    switch (Result) {
      case 0xF6:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Parameter Not Readable\n"));
        break;
      case 0xF7:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Parameter Not Writeable\n"));
        break;
      case 0xF8:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Parameter Already Written\n"));
        break;
      case 0xF9:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Invalid Length\n"));
        break;
      case 0xFA:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Invalid Value\n"));
        break;
      case 0xFB:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Invalid Selector\n"));
        break;
      case 0xFC:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Invalid Index\n"));
        break;
      case 0xFD:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Invalid Idn\n"));
        break;
      case 0xFE:
        DEBUG ((DEBUG_VERBOSE, "Query Response with Invalid Opcode\n"));
        break;
      case 0xFF:
        DEBUG ((DEBUG_VERBOSE, "Query Response with General Failure\n"));
        break;
      default:
        ASSERT (FALSE);
        break;
    }
  }
  
  /**
    Swap little endian to big endian.
  
    @param[in, out] Buffer      The data buffer. In input, it contains little endian data.
                                In output, it will become big endian.
    @param[in]      BufferSize  The length of converted data.
  
  **/
  VOID
  SwapLittleEndianToBigEndian (
    IN OUT UINT8   *Buffer,
    IN     UINT32  BufferSize
    )
  {
    UINT32  Index;
    UINT8   Temp;
    UINT32  SwapCount;
  
    SwapCount = BufferSize / 2;
    for (Index = 0; Index < SwapCount; Index++) {
      Temp                           = Buffer[Index];
      Buffer[Index]                  = Buffer[BufferSize - 1 - Index];
      Buffer[BufferSize - 1 - Index] = Temp;
    }
  }
  
  /**
    Fill TSF field of QUERY REQUEST UPIU.
  
    @param[in, out] TsfBase      The base address of TSF field of QUERY REQUEST UPIU.
    @param[in]      Opcode       The opcode of request.
    @param[in]      DescId       The descriptor ID of request.
    @param[in]      Index        The index of request.
    @param[in]      Selector     The selector of request.
    @param[in]      Length       The length of transferred data. The maximum is 4.
    @param[in]      Value        The value of transferred data.
  
  **/
  VOID
  UfsFillTsfOfQueryReqUpiu (
    IN OUT UTP_UPIU_TSF  *TsfBase,
    IN     UINT8         Opcode,
    IN     UINT8         DescId    OPTIONAL,
    IN     UINT8         Index     OPTIONAL,
    IN     UINT8         Selector  OPTIONAL,
    IN     UINT16        Length    OPTIONAL,
    IN     UINT32        Value     OPTIONAL
    )
  {
    ASSERT (TsfBase != NULL);
    ASSERT (Opcode <= UtpQueryFuncOpcodeTogFlag);
  
    TsfBase->Opcode = Opcode;
    if (Opcode != UtpQueryFuncOpcodeNop) {
      TsfBase->DescId   = DescId;
      TsfBase->Index    = Index;
      TsfBase->Selector = Selector;
  
      if ((Opcode == UtpQueryFuncOpcodeRdDesc) || (Opcode == UtpQueryFuncOpcodeWrDesc)) {
        SwapLittleEndianToBigEndian ((UINT8 *)&Length, sizeof (Length));
        TsfBase->Length = Length;
      }
  
      if (Opcode == UtpQueryFuncOpcodeWrAttr) {
        SwapLittleEndianToBigEndian ((UINT8 *)&Value, sizeof (Value));
        TsfBase->Value = Value;
      }
    }
  }
  
  /**
    Initialize COMMAND UPIU.
  
    @param[in, out] Command         The base address of COMMAND UPIU.
    @param[in]      Lun             The Lun on which the SCSI command is executed.
    @param[in]      TaskTag         The task tag of request.
    @param[in]      Cdb             The cdb buffer containing SCSI command.
    @param[in]      CdbLength       The cdb length.
    @param[in]      DataDirection   The direction of data transfer.
    @param[in]      ExpDataTranLen  The expected transfer data length.
  
    @retval EFI_SUCCESS     The initialization succeed.
  
  **/
  EFI_STATUS
  UfsInitCommandUpiu (
    IN OUT UTP_COMMAND_UPIU    *Command,
    IN     UINT8               Lun,
    IN     UINT8               TaskTag,
    IN     UINT8               *Cdb,
    IN     UINT8               CdbLength,
    IN     UFS_DATA_DIRECTION  DataDirection,
    IN     UINT32              ExpDataTranLen
    )
  {
    UINT8  Flags;
  
    ASSERT ((Command != NULL) && (Cdb != NULL));
  
    //
    // Task attribute is hard-coded to Ordered.
    //
    if (DataDirection == UfsDataIn) {
      Flags = BIT0 | BIT6;
    } else if (DataDirection == UfsDataOut) {
      Flags = BIT0 | BIT5;
    } else {
      Flags = BIT0;
    }
  
    //
    // Fill UTP COMMAND UPIU associated fields.
    //
    Command->TransCode = 0x01;
    Command->Flags     = Flags;
    Command->Lun       = Lun;
    Command->TaskTag   = TaskTag;
    Command->CmdSet    = 0x00;
    SwapLittleEndianToBigEndian ((UINT8 *)&ExpDataTranLen, sizeof (ExpDataTranLen));
    Command->ExpDataTranLen = ExpDataTranLen;
  
    CopyMem (Command->Cdb, Cdb, CdbLength);
  
    return EFI_SUCCESS;
  }
  
  /**
    Initialize UTP PRDT for data transfer.
  
    @param[in] Prdt         The base address of PRDT.
    @param[in] Buffer       The buffer to be read or written.
    @param[in] BufferSize   The data size to be read or written.
  
    @retval EFI_SUCCESS     The initialization succeed.
  
  **/
  EFI_STATUS
  UfsInitUtpPrdt (
    IN  UTP_TR_PRD  *Prdt,
    IN  VOID        *Buffer,
    IN  UINT32      BufferSize
    )
  {
    UINT32  PrdtIndex;
    UINT32  RemainingLen;
    UINT8   *Remaining;
    UINTN   PrdtNumber;
  
    if ((BufferSize & (BIT0 | BIT1)) != 0) {
      BufferSize &= ~(BIT0 | BIT1);
      DEBUG ((DEBUG_WARN, "UfsInitUtpPrdt: The BufferSize [%d] is not dword-aligned!\n", BufferSize));
    }
  
    if (BufferSize == 0) {
      return EFI_SUCCESS;
    }
  
    ASSERT (((UINTN)Buffer & (BIT0 | BIT1)) == 0);
  
    RemainingLen = BufferSize;
    Remaining    = Buffer;
    PrdtNumber   = (UINTN)DivU64x32 ((UINT64)BufferSize + UFS_MAX_DATA_LEN_PER_PRD - 1, UFS_MAX_DATA_LEN_PER_PRD);
  
    for (PrdtIndex = 0; PrdtIndex < PrdtNumber; PrdtIndex++) {
      if (RemainingLen < UFS_MAX_DATA_LEN_PER_PRD) {
        Prdt[PrdtIndex].DbCount = (UINT32)RemainingLen - 1;
      } else {
        Prdt[PrdtIndex].DbCount = UFS_MAX_DATA_LEN_PER_PRD - 1;
      }
  
      Prdt[PrdtIndex].DbAddr  = (UINT32)RShiftU64 ((UINT64)(UINTN)Remaining, 2);
      Prdt[PrdtIndex].DbAddrU = (UINT32)RShiftU64 ((UINT64)(UINTN)Remaining, 32);
      RemainingLen           -= UFS_MAX_DATA_LEN_PER_PRD;
      Remaining              += UFS_MAX_DATA_LEN_PER_PRD;
    }
  
    return EFI_SUCCESS;
  }
  
  /**
    Initialize QUERY REQUEST UPIU.
  
    @param[in, out] QueryReq      The base address of QUERY REQUEST UPIU.
    @param[in]      TaskTag       The task tag of request.
    @param[in]      Opcode        The opcode of request.
    @param[in]      DescId        The descriptor ID of request.
    @param[in]      Index         The index of request.
    @param[in]      Selector      The selector of request.
    @param[in]      DataSize      The data size to be read or written.
    @param[in]      Data          The buffer to be read or written.
  
    @retval EFI_SUCCESS           The initialization succeed.
  
  **/
  EFI_STATUS
  UfsInitQueryRequestUpiu (
    IN OUT UTP_QUERY_REQ_UPIU  *QueryReq,
    IN     UINT8               TaskTag,
    IN     UINT8               Opcode,
    IN     UINT8               DescId,
    IN     UINT8               Index,
    IN     UINT8               Selector,
    IN     UINTN               DataSize   OPTIONAL,
    IN     UINT8               *Data      OPTIONAL
    )
  {
    ASSERT (QueryReq != NULL);
  
    QueryReq->TransCode = 0x16;
    QueryReq->TaskTag   = TaskTag;
    if ((Opcode == UtpQueryFuncOpcodeRdDesc) || (Opcode == UtpQueryFuncOpcodeRdFlag) || (Opcode == UtpQueryFuncOpcodeRdAttr)) {
      QueryReq->QueryFunc = QUERY_FUNC_STD_READ_REQ;
    } else {
      QueryReq->QueryFunc = QUERY_FUNC_STD_WRITE_REQ;
    }
  
    if (Opcode == UtpQueryFuncOpcodeWrAttr) {
      UfsFillTsfOfQueryReqUpiu (&QueryReq->Tsf, Opcode, DescId, Index, Selector, 0, *(UINT32 *)Data);
    } else if ((Opcode == UtpQueryFuncOpcodeRdDesc) || (Opcode == UtpQueryFuncOpcodeWrDesc)) {
      UfsFillTsfOfQueryReqUpiu (&QueryReq->Tsf, Opcode, DescId, Index, Selector, (UINT16)DataSize, 0);
    } else {
      UfsFillTsfOfQueryReqUpiu (&QueryReq->Tsf, Opcode, DescId, Index, Selector, 0, 0);
    }
  
    if (Opcode == UtpQueryFuncOpcodeWrDesc) {
      CopyMem (QueryReq + 1, Data, DataSize);
  
      SwapLittleEndianToBigEndian ((UINT8 *)&DataSize, sizeof (UINT16));
      QueryReq->DataSegLen = (UINT16)DataSize;
    }
  
    return EFI_SUCCESS;
  }
  
  /**
    Allocate COMMAND/RESPONSE UPIU for filling UTP TRD's command descriptor field.
  
    @param[in]  Private           The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]  Lun               The Lun on which the SCSI command is executed.
    @param[in]  Packet            The pointer to the UFS_SCSI_REQUEST_PACKET data structure.
    @param[in]  Trd               The pointer to the UTP Transfer Request Descriptor.
    @param[out] BufferMap         A resulting value, if not NULL, to pass to IoMmuUnmap().
  
    @retval EFI_SUCCESS           The creation succeed.
    @retval EFI_DEVICE_ERROR      The creation failed.
    @retval EFI_OUT_OF_RESOURCES  The memory resource is insufficient.
  
  **/
  EFI_STATUS
  UfsCreateScsiCommandDesc (
    IN     UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN     UINT8                     Lun,
    IN     UFS_SCSI_REQUEST_PACKET   *Packet,
    IN     UTP_TRD                   *Trd,
    OUT VOID                         **BufferMap
    )
  {
    UINT8                  *CommandDesc;
    UINTN                  TotalLen;
    UINTN                  PrdtNumber;
    VOID                   *Buffer;
    UINT32                 Length;
    UTP_COMMAND_UPIU       *CommandUpiu;
    UTP_TR_PRD             *PrdtBase;
    UFS_DATA_DIRECTION     DataDirection;
    EFI_STATUS             Status;
    EDKII_IOMMU_OPERATION  MapOp;
    UINTN                  MapLength;
    EFI_PHYSICAL_ADDRESS   BufferPhyAddr;
  
    ASSERT ((Private != NULL) && (Packet != NULL) && (Trd != NULL));
  
    BufferPhyAddr = 0;
  
    if (Packet->DataDirection == UfsDataIn) {
      Buffer        = Packet->InDataBuffer;
      Length        = Packet->InTransferLength;
      DataDirection = UfsDataIn;
      MapOp         = EdkiiIoMmuOperationBusMasterWrite;
    } else {
      Buffer        = Packet->OutDataBuffer;
      Length        = Packet->OutTransferLength;
      DataDirection = UfsDataOut;
      MapOp         = EdkiiIoMmuOperationBusMasterRead;
    }
  
    if (Length == 0) {
      DataDirection = UfsNoData;
    } else {
      MapLength = Length;
      Status    = IoMmuMap (MapOp, Buffer, &MapLength, &BufferPhyAddr, BufferMap);
  
      if (EFI_ERROR (Status) || (MapLength != Length)) {
        DEBUG ((DEBUG_ERROR, "UfsCreateScsiCommandDesc: Fail to map data buffer.\n"));
        return EFI_OUT_OF_RESOURCES;
      }
    }
  
    PrdtNumber = (UINTN)DivU64x32 ((UINT64)Length + UFS_MAX_DATA_LEN_PER_PRD - 1, UFS_MAX_DATA_LEN_PER_PRD);
  
    TotalLen    = ROUNDUP8 (sizeof (UTP_COMMAND_UPIU)) + ROUNDUP8 (sizeof (UTP_RESPONSE_UPIU)) + PrdtNumber * sizeof (UTP_TR_PRD);
    CommandDesc = UfsPeimAllocateMem (Private->Pool, TotalLen);
    if (CommandDesc == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }
  
    CommandUpiu = (UTP_COMMAND_UPIU *)CommandDesc;
    PrdtBase    = (UTP_TR_PRD *)(CommandDesc + ROUNDUP8 (sizeof (UTP_COMMAND_UPIU)) + ROUNDUP8 (sizeof (UTP_RESPONSE_UPIU)));
  
    UfsInitCommandUpiu (CommandUpiu, Lun, Private->TaskTag++, Packet->Cdb, Packet->CdbLength, DataDirection, Length);
    UfsInitUtpPrdt (PrdtBase, (VOID *)(UINTN)BufferPhyAddr, Length);
  
    //
    // Fill UTP_TRD associated fields
    // NOTE: Some UFS host controllers request the Response UPIU and the Physical Region Description Table
    // *MUST* be located at a 64-bit aligned boundary.
    //
    Trd->Int    = UFS_INTERRUPT_COMMAND;
    Trd->Dd     = DataDirection;
    Trd->Ct     = UFS_STORAGE_COMMAND_TYPE;
    Trd->Ocs    = UFS_HC_TRD_OCS_INIT_VALUE;
    Trd->UcdBa  = (UINT32)RShiftU64 ((UINT64)(UINTN)CommandUpiu, 7);
    Trd->UcdBaU = (UINT32)RShiftU64 ((UINT64)(UINTN)CommandUpiu, 32);
    Trd->RuL    = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_RESPONSE_UPIU)), sizeof (UINT32));
    Trd->RuO    = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_COMMAND_UPIU)), sizeof (UINT32));
    Trd->PrdtL  = (UINT16)PrdtNumber;
    Trd->PrdtO  = (UINT16)DivU64x32 ((UINT64)(ROUNDUP8 (sizeof (UTP_COMMAND_UPIU)) + ROUNDUP8 (sizeof (UTP_RESPONSE_UPIU))), sizeof (UINT32));
    return EFI_SUCCESS;
  }
  
  /**
    Allocate QUERY REQUEST/QUERY RESPONSE UPIU for filling UTP TRD's command descriptor field.
  
    @param[in]  Private           The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]  Packet            The pointer to the UFS_DEVICE_MANAGEMENT_REQUEST_PACKET data structure.
    @param[in]  Trd               The pointer to the UTP Transfer Request Descriptor.
  
    @retval EFI_SUCCESS           The creation succeed.
    @retval EFI_DEVICE_ERROR      The creation failed.
    @retval EFI_OUT_OF_RESOURCES  The memory resource is insufficient.
    @retval EFI_INVALID_PARAMETER The parameter passed in is invalid.
  
  **/
  EFI_STATUS
  UfsCreateDMCommandDesc (
    IN  UFS_PEIM_HC_PRIVATE_DATA              *Private,
    IN  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet,
    IN  UTP_TRD                               *Trd
    )
  {
    UINT8               *CommandDesc;
    UINTN               TotalLen;
    UTP_QUERY_REQ_UPIU  *QueryReqUpiu;
    UINT8               Opcode;
    UINT32              DataSize;
    UINT8               *Data;
    UINT8               DataDirection;
  
    ASSERT ((Private != NULL) && (Packet != NULL) && (Trd != NULL));
  
    Opcode = Packet->Opcode;
    if ((Opcode > UtpQueryFuncOpcodeTogFlag) || (Opcode == UtpQueryFuncOpcodeNop)) {
      return EFI_INVALID_PARAMETER;
    }
  
    DataDirection = Packet->DataDirection;
    if (DataDirection == UfsDataIn) {
      DataSize = Packet->InTransferLength;
      Data     = Packet->InDataBuffer;
    } else if (DataDirection == UfsDataOut) {
      DataSize = Packet->OutTransferLength;
      Data     = Packet->OutDataBuffer;
    } else {
      DataSize = 0;
      Data     = NULL;
    }
  
    if (((Opcode != UtpQueryFuncOpcodeRdFlag) && (Opcode != UtpQueryFuncOpcodeSetFlag) &&
         (Opcode != UtpQueryFuncOpcodeClrFlag) && (Opcode != UtpQueryFuncOpcodeTogFlag)) &&
        ((DataSize == 0) || (Data == NULL)))
    {
      return EFI_INVALID_PARAMETER;
    }
  
    if ((Opcode == UtpQueryFuncOpcodeWrAttr) && (DataSize != sizeof (UINT32))) {
      return EFI_INVALID_PARAMETER;
    }
  
    if ((Opcode == UtpQueryFuncOpcodeWrDesc) || (Opcode == UtpQueryFuncOpcodeRdDesc)) {
      TotalLen = ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)) + ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU)) + ROUNDUP8 (DataSize);
    } else {
      TotalLen = ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)) + ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU));
    }
  
    CommandDesc = UfsPeimAllocateMem (Private->Pool, TotalLen);
    if (CommandDesc == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }
  
    //
    // Initialize UTP QUERY REQUEST UPIU
    //
    QueryReqUpiu = (UTP_QUERY_REQ_UPIU *)CommandDesc;
    UfsInitQueryRequestUpiu (
      QueryReqUpiu,
      Private->TaskTag++,
      Opcode,
      Packet->DescId,
      Packet->Index,
      Packet->Selector,
      DataSize,
      Data
      );
  
    //
    // Fill UTP_TRD associated fields
    // NOTE: Some UFS host controllers request the Query Response UPIU *MUST* be located at a 64-bit aligned boundary.
    //
    Trd->Int    = UFS_INTERRUPT_COMMAND;
    Trd->Dd     = DataDirection;
    Trd->Ct     = UFS_STORAGE_COMMAND_TYPE;
    Trd->Ocs    = UFS_HC_TRD_OCS_INIT_VALUE;
    Trd->UcdBa  = (UINT32)RShiftU64 ((UINT64)(UINTN)QueryReqUpiu, 7);
    Trd->UcdBaU = (UINT32)RShiftU64 ((UINT64)(UINTN)QueryReqUpiu, 32);
    if (Opcode == UtpQueryFuncOpcodeWrDesc) {
      Trd->RuL = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU)), sizeof (UINT32));
      Trd->RuO = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)) + ROUNDUP8 (DataSize), sizeof (UINT32));
    } else {
      Trd->RuL = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_RESP_UPIU)) + ROUNDUP8 (DataSize), sizeof (UINT32));
      Trd->RuO = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_QUERY_REQ_UPIU)), sizeof (UINT32));
    }
  
    return EFI_SUCCESS;
  }
  
  /**
    Allocate NOP IN and NOP OUT UPIU for filling UTP TRD's command descriptor field.
  
    @param[in]  Private           The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]  Trd               The pointer to the UTP Transfer Request Descriptor.
  
    @retval EFI_SUCCESS           The creation succeed.
    @retval EFI_DEVICE_ERROR      The creation failed.
    @retval EFI_OUT_OF_RESOURCES  The memory resource is insufficient.
  
  **/
  EFI_STATUS
  UfsCreateNopCommandDesc (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN  UTP_TRD                   *Trd
    )
  {
    UINT8             *CommandDesc;
    UINTN             TotalLen;
    UTP_NOP_OUT_UPIU  *NopOutUpiu;
  
    ASSERT ((Private != NULL) && (Trd != NULL));
  
    TotalLen    = ROUNDUP8 (sizeof (UTP_NOP_OUT_UPIU)) + ROUNDUP8 (sizeof (UTP_NOP_IN_UPIU));
    CommandDesc = UfsPeimAllocateMem (Private->Pool, TotalLen);
    if (CommandDesc == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }
  
    NopOutUpiu = (UTP_NOP_OUT_UPIU *)CommandDesc;
  
    NopOutUpiu->TaskTag = Private->TaskTag++;
  
    //
    // Fill UTP_TRD associated fields
    // NOTE: Some UFS host controllers request the Nop Out UPIU *MUST* be located at a 64-bit aligned boundary.
    //
    Trd->Int    = UFS_INTERRUPT_COMMAND;
    Trd->Dd     = 0x00;
    Trd->Ct     = UFS_STORAGE_COMMAND_TYPE;
    Trd->Ocs    = UFS_HC_TRD_OCS_INIT_VALUE;
    Trd->UcdBa  = (UINT32)RShiftU64 ((UINT64)(UINTN)NopOutUpiu, 7);
    Trd->UcdBaU = (UINT32)RShiftU64 ((UINT64)(UINTN)NopOutUpiu, 32);
    Trd->RuL    = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_NOP_IN_UPIU)), sizeof (UINT32));
    Trd->RuO    = (UINT16)DivU64x32 ((UINT64)ROUNDUP8 (sizeof (UTP_NOP_OUT_UPIU)), sizeof (UINT32));
  
    return EFI_SUCCESS;
  }
  
  /**
    Find out available slot in transfer list of a UFS device.
  
    @param[in]  Private       The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[out] Slot          The available slot.
  
    @retval EFI_SUCCESS       The available slot was found successfully.
  
  **/
  EFI_STATUS
  UfsFindAvailableSlotInTrl (
    IN     UFS_PEIM_HC_PRIVATE_DATA  *Private,
    OUT UINT8                        *Slot
    )
  {
    ASSERT ((Private != NULL) && (Slot != NULL));
  
    //
    // The simplest algo to always use slot 0.
    // TODO: enhance it to support async transfer with multiple slot.
    //
    *Slot = 0;
  
    return EFI_SUCCESS;
  }
  
  /**
    Start specified slot in transfer list of a UFS device.
  
    @param[in]  Private       The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]  Slot          The slot to be started.
  
  **/
  VOID
  UfsStartExecCmd (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN  UINT8                     Slot
    )
  {
    UINTN   UfsHcBase;
    UINTN   Address;
    UINT32  Data;
  
    UfsHcBase = Private->UfsHcBase;
  
    Address = UfsHcBase + UFS_HC_UTRLRSR_OFFSET;
    Data    = MmioRead32 (Address);
    if ((Data & UFS_HC_UTRLRSR) != UFS_HC_UTRLRSR) {
      MmioWrite32 (Address, UFS_HC_UTRLRSR);
    }
  
    Address = UfsHcBase + UFS_HC_UTRLDBR_OFFSET;
    MmioWrite32 (Address, BIT0 << Slot);
  }
  
  /**
    Stop specified slot in transfer list of a UFS device.
  
    @param[in]  Private       The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]  Slot          The slot to be stop.
  
  **/
  VOID
  UfsStopExecCmd (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN  UINT8                     Slot
    )
  {
    UINTN   UfsHcBase;
    UINTN   Address;
    UINT32  Data;
  
    UfsHcBase = Private->UfsHcBase;
  
    Address = UfsHcBase + UFS_HC_UTRLDBR_OFFSET;
    Data    = MmioRead32 (Address);
    if ((Data & (BIT0 << Slot)) != 0) {
      Address = UfsHcBase + UFS_HC_UTRLCLR_OFFSET;
      Data    = MmioRead32 (Address);
      MmioWrite32 (Address, (Data & ~(BIT0 << Slot)));
    }
  }
  
  /**
    Extracts return data from query response upiu.
  
    @param[in, out] Packet        Pointer to the UFS_DEVICE_MANAGEMENT_REQUEST_PACKET.
    @param[in]      QueryResp     Pointer to the query response.
  
    @retval EFI_INVALID_PARAMETER Packet or QueryResp are empty or opcode is invalid.
    @retval EFI_DEVICE_ERROR      Data returned from device is invalid.
    @retval EFI_SUCCESS           Data extracted.
  
  **/
  EFI_STATUS
  UfsGetReturnDataFromQueryResponse (
    IN OUT UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet,
    IN     UTP_QUERY_RESP_UPIU                   *QueryResp
    )
  {
    UINT16  ReturnDataSize;
  
    ReturnDataSize = 0;
  
    if ((Packet == NULL) || (QueryResp == NULL)) {
      return EFI_INVALID_PARAMETER;
    }
  
    switch (Packet->Opcode) {
      case UtpQueryFuncOpcodeRdDesc:
        ReturnDataSize = QueryResp->Tsf.Length;
        SwapLittleEndianToBigEndian ((UINT8 *)&ReturnDataSize, sizeof (UINT16));
        //
        // Make sure the hardware device does not return more data than expected.
        //
        if (ReturnDataSize > Packet->InTransferLength) {
          return EFI_DEVICE_ERROR;
        }
  
        CopyMem (Packet->InDataBuffer, (QueryResp + 1), ReturnDataSize);
        Packet->InTransferLength = ReturnDataSize;
        break;
      case UtpQueryFuncOpcodeWrDesc:
        ReturnDataSize = QueryResp->Tsf.Length;
        SwapLittleEndianToBigEndian ((UINT8 *)&ReturnDataSize, sizeof (UINT16));
        Packet->OutTransferLength = ReturnDataSize;
        break;
      case UtpQueryFuncOpcodeRdFlag:
        //
        // The 'FLAG VALUE' field is at byte offset 3 of QueryResp->Tsf.Value
        //
        *((UINT8 *)(Packet->InDataBuffer)) = *((UINT8 *)&(QueryResp->Tsf.Value) + 3);
        break;
      case UtpQueryFuncOpcodeSetFlag:
      case UtpQueryFuncOpcodeClrFlag:
      case UtpQueryFuncOpcodeTogFlag:
        //
        // The 'FLAG VALUE' field is at byte offset 3 of QueryResp->Tsf.Value
        //
        *((UINT8 *)(Packet->OutDataBuffer)) = *((UINT8 *)&(QueryResp->Tsf.Value) + 3);
        break;
      default:
        return EFI_INVALID_PARAMETER;
    }
  
    return EFI_SUCCESS;
  }
  
  /**
    Creates Transfer Request descriptor and sends Query Request to the device.
  
    @param[in]      Private       Pointer to the UFS_PEIM_HC_PRIVATE_DATA.
    @param[in, out] Packet        Pointer to the UFS_DEVICE_MANAGEMENT_REQUEST_PACKET.
  
    @retval EFI_SUCCESS           The device descriptor was read/written successfully.
    @retval EFI_INVALID_PARAMETER The DescId, Index and Selector fields in Packet are invalid
                                  combination to point to a type of UFS device descriptor.
    @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the device descriptor.
    @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the device descriptor.
  
  **/
  EFI_STATUS
  UfsSendDmRequestRetry (
    IN     UFS_PEIM_HC_PRIVATE_DATA              *Private,
    IN OUT UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet
    )
  {
    UINT8                Slot;
    EFI_STATUS           Status;
    UTP_TRD              *Trd;
    UINTN                Address;
    UTP_QUERY_RESP_UPIU  *QueryResp;
    UINT8                *CmdDescBase;
    UINT32               CmdDescSize;
  
    //
    // Find out which slot of transfer request list is available.
    //
    Status = UfsFindAvailableSlotInTrl (Private, &Slot);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  
    Trd = ((UTP_TRD *)Private->UtpTrlBase) + Slot;
    //
    // Fill transfer request descriptor to this slot.
    //
    Status = UfsCreateDMCommandDesc (Private, Packet, Trd);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "Failed to create DM command descriptor\n"));
      return Status;
    }
  
    //
    // Check the transfer request result.
    //
    CmdDescBase = (UINT8 *)(UINTN)(LShiftU64 ((UINT64)Trd->UcdBaU, 32) | LShiftU64 ((UINT64)Trd->UcdBa, 7));
    QueryResp   = (UTP_QUERY_RESP_UPIU *)(CmdDescBase + Trd->RuO * sizeof (UINT32));
    CmdDescSize = Trd->RuO * sizeof (UINT32) + Trd->RuL * sizeof (UINT32);
  
    //
    // Start to execute the transfer request.
    //
    UfsStartExecCmd (Private, Slot);
  
    //
    // Wait for the completion of the transfer request.
    //
    Address = Private->UfsHcBase + UFS_HC_UTRLDBR_OFFSET;
    Status  = UfsWaitMemSet (Address, (BIT0 << Slot), 0, Packet->Timeout);
    if (EFI_ERROR (Status)) {
      goto Exit;
    }
  
    if ((Trd->Ocs != 0) || (QueryResp->QueryResp != UfsUtpQueryResponseSuccess)) {
      DEBUG ((DEBUG_ERROR, "Failed to send query request, OCS = %X, QueryResp = %X\n", Trd->Ocs, QueryResp->QueryResp));
      DumpQueryResponseResult (QueryResp->QueryResp);
      Status = EFI_DEVICE_ERROR;
      goto Exit;
    }
  
    Status = UfsGetReturnDataFromQueryResponse (Packet, QueryResp);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "Failed to get return data from query response\n"));
      goto Exit;
    }
  
  Exit:
    UfsStopExecCmd (Private, Slot);
    UfsPeimFreeMem (Private->Pool, CmdDescBase, CmdDescSize);
  
    return Status;
  }
  
  /**
    Sends Query Request to the device. Query is sent until device responds correctly or counter runs out.
  
    @param[in]      Private       Pointer to the UFS_PEIM_HC_PRIVATE_DATA.
    @param[in, out] Packet        Pointer to the UFS_DEVICE_MANAGEMENT_REQUEST_PACKET.
  
    @retval EFI_SUCCESS           The device responded correctly to the Query request.
    @retval EFI_INVALID_PARAMETER The DescId, Index and Selector fields in Packet are invalid
                                  combination to point to a type of UFS device descriptor.
    @retval EFI_DEVICE_ERROR      A device error occurred while waiting for the response from the device.
    @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of the operation.
  
  **/
  EFI_STATUS
  UfsSendDmRequest (
    IN     UFS_PEIM_HC_PRIVATE_DATA              *Private,
    IN OUT UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet
    )
  {
    EFI_STATUS  Status;
    UINT8       Retry;
  
    Status = EFI_SUCCESS;
  
    for (Retry = 0; Retry < 5; Retry++) {
      Status = UfsSendDmRequestRetry (Private, Packet);
      if (!EFI_ERROR (Status)) {
        return EFI_SUCCESS;
      }
    }
  
    DEBUG ((DEBUG_ERROR, "Failed to get response from the device after %d retries\n", Retry));
    return Status;
  }
  
  /**
    Read or write specified device descriptor of a UFS device.
  
    @param[in]      Private       The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]      Read          The boolean variable to show r/w direction.
    @param[in]      DescId        The ID of device descriptor.
    @param[in]      Index         The Index of device descriptor.
    @param[in]      Selector      The Selector of device descriptor.
    @param[in, out] Descriptor    The buffer of device descriptor to be read or written.
    @param[in]      DescSize      The size of device descriptor buffer.
  
    @retval EFI_SUCCESS           The device descriptor was read/written successfully.
    @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the device descriptor.
    @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the device descriptor.
  
  **/
  EFI_STATUS
  UfsRwDeviceDesc (
    IN     UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN     BOOLEAN                   Read,
    IN     UINT8                     DescId,
    IN     UINT8                     Index,
    IN     UINT8                     Selector,
    IN OUT VOID                      *Descriptor,
    IN     UINT32                    DescSize
    )
  {
    EFI_STATUS                            Status;
    UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  Packet;
  
    ZeroMem (&Packet, sizeof (UFS_DEVICE_MANAGEMENT_REQUEST_PACKET));
  
    if (Read) {
      Packet.DataDirection    = UfsDataIn;
      Packet.InDataBuffer     = Descriptor;
      Packet.InTransferLength = DescSize;
      Packet.Opcode           = UtpQueryFuncOpcodeRdDesc;
    } else {
      Packet.DataDirection     = UfsDataOut;
      Packet.OutDataBuffer     = Descriptor;
      Packet.OutTransferLength = DescSize;
      Packet.Opcode            = UtpQueryFuncOpcodeWrDesc;
    }
  
    Packet.DescId   = DescId;
    Packet.Index    = Index;
    Packet.Selector = Selector;
    Packet.Timeout  = UFS_TIMEOUT;
  
    Status = UfsSendDmRequest (Private, &Packet);
    return Status;
  }
  
  /**
    Read or write specified flag of a UFS device.
  
    @param[in]      Private       The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]      Read          The boolean variable to show r/w direction.
    @param[in]      FlagId        The ID of flag to be read or written.
    @param[in, out] Value         The value to set or clear flag.
  
    @retval EFI_SUCCESS           The flag was read/written successfully.
    @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the flag.
    @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the flag.
  
  **/
  EFI_STATUS
  UfsRwFlags (
    IN     UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN     BOOLEAN                   Read,
    IN     UINT8                     FlagId,
    IN OUT UINT8                     *Value
    )
  {
    EFI_STATUS                            Status;
    UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  Packet;
  
    if (Value == NULL) {
      return EFI_INVALID_PARAMETER;
    }
  
    ZeroMem (&Packet, sizeof (UFS_DEVICE_MANAGEMENT_REQUEST_PACKET));
  
    if (Read) {
      ASSERT (Value != NULL);
      Packet.DataDirection    = UfsDataIn;
      Packet.InDataBuffer     = (VOID *)Value;
      Packet.InTransferLength = 0;
      Packet.Opcode           = UtpQueryFuncOpcodeRdFlag;
    } else {
      Packet.DataDirection     = UfsDataOut;
      Packet.OutDataBuffer     = (VOID *)Value;
      Packet.OutTransferLength = 0;
      if (*Value == 1) {
        Packet.Opcode = UtpQueryFuncOpcodeSetFlag;
      } else if (*Value == 0) {
        Packet.Opcode = UtpQueryFuncOpcodeClrFlag;
      } else {
        return EFI_INVALID_PARAMETER;
      }
    }
  
    Packet.DescId   = FlagId;
    Packet.Index    = 0;
    Packet.Selector = 0;
    Packet.Timeout  = UFS_TIMEOUT;
  
    Status = UfsSendDmRequest (Private, &Packet);
  
    return Status;
  }
  
  /**
    Set specified flag to 1 on a UFS device.
  
    @param[in]  Private           The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]  FlagId            The ID of flag to be set.
  
    @retval EFI_SUCCESS           The flag was set successfully.
    @retval EFI_DEVICE_ERROR      A device error occurred while attempting to set the flag.
    @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of setting the flag.
  
  **/
  EFI_STATUS
  UfsSetFlag (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN  UINT8                     FlagId
    )
  {
    EFI_STATUS  Status;
    UINT8       Value;
  
    Value  = 1;
    Status = UfsRwFlags (Private, FALSE, FlagId, &Value);
  
    return Status;
  }
  
  /**
    Sends NOP IN cmd to a UFS device for initialization process request.
    For more details, please refer to UFS 2.0 spec Figure 13.3.
  
    @param[in]  Private           The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
  
    @retval EFI_SUCCESS           The NOP IN command was sent by the host. The NOP OUT response was
                                  received successfully.
    @retval EFI_DEVICE_ERROR      A device error occurred while attempting to execute NOP IN command.
    @retval EFI_OUT_OF_RESOURCES  The resource for transfer is not available.
    @retval EFI_TIMEOUT           A timeout occurred while waiting for the NOP IN command to execute.
  
  **/
  EFI_STATUS
  UfsExecNopCmds (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private
    )
  {
    EFI_STATUS       Status;
    UINT8            Slot;
    UTP_TRD          *Trd;
    UTP_NOP_IN_UPIU  *NopInUpiu;
    UINT8            *CmdDescBase;
    UINT32           CmdDescSize;
    UINTN            Address;
  
    //
    // Find out which slot of transfer request list is available.
    //
    Status = UfsFindAvailableSlotInTrl (Private, &Slot);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  
    Trd    = ((UTP_TRD *)Private->UtpTrlBase) + Slot;
    Status = UfsCreateNopCommandDesc (Private, Trd);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  
    //
    // Check the transfer request result.
    //
    CmdDescBase = (UINT8 *)(UINTN)(LShiftU64 ((UINT64)Trd->UcdBaU, 32) | LShiftU64 ((UINT64)Trd->UcdBa, 7));
    NopInUpiu   = (UTP_NOP_IN_UPIU *)(CmdDescBase + Trd->RuO * sizeof (UINT32));
    CmdDescSize = Trd->RuO * sizeof (UINT32) + Trd->RuL * sizeof (UINT32);
  
    //
    // Start to execute the transfer request.
    //
    UfsStartExecCmd (Private, Slot);
  
    //
    // Wait for the completion of the transfer request.
    //
    Address = Private->UfsHcBase + UFS_HC_UTRLDBR_OFFSET;
    Status  = UfsWaitMemSet (Address, BIT0 << Slot, 0, UFS_TIMEOUT);
    if (EFI_ERROR (Status)) {
      goto Exit;
    }
  
    if (NopInUpiu->Resp != 0) {
      Status = EFI_DEVICE_ERROR;
    } else {
      Status = EFI_SUCCESS;
    }
  
  Exit:
    UfsStopExecCmd (Private, Slot);
    UfsPeimFreeMem (Private->Pool, CmdDescBase, CmdDescSize);
  
    return Status;
  }
  
  /**
    Sends a UFS-supported SCSI Request Packet to a UFS device that is attached to the UFS host controller.
  
    @param[in]      Private       The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in]      Lun           The LUN of the UFS device to send the SCSI Request Packet.
    @param[in, out] Packet        A pointer to the SCSI Request Packet to send to a specified Lun of the
                                  UFS device.
  
    @retval EFI_SUCCESS           The SCSI Request Packet was sent by the host. For bi-directional
                                  commands, InTransferLength bytes were transferred from
                                  InDataBuffer. For write and bi-directional commands,
                                  OutTransferLength bytes were transferred by
                                  OutDataBuffer.
    @retval EFI_DEVICE_ERROR      A device error occurred while attempting to send the SCSI Request
                                  Packet.
    @retval EFI_OUT_OF_RESOURCES  The resource for transfer is not available.
    @retval EFI_TIMEOUT           A timeout occurred while waiting for the SCSI Request Packet to execute.
  
  **/
  EFI_STATUS
  UfsExecScsiCmds (
    IN     UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN     UINT8                     Lun,
    IN OUT UFS_SCSI_REQUEST_PACKET   *Packet
    )
  {
    EFI_STATUS         Status;
    UINT8              Slot;
    UTP_TRD            *Trd;
    UINTN              Address;
    UINT8              *CmdDescBase;
    UINT32             CmdDescSize;
    UTP_RESPONSE_UPIU  *Response;
    UINT16             SenseDataLen;
    UINT32             ResTranCount;
    VOID               *PacketBufferMap;
  
    //
    // Find out which slot of transfer request list is available.
    //
    Status = UfsFindAvailableSlotInTrl (Private, &Slot);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  
    Trd             = ((UTP_TRD *)Private->UtpTrlBase) + Slot;
    PacketBufferMap = NULL;
  
    //
    // Fill transfer request descriptor to this slot.
    //
    Status = UfsCreateScsiCommandDesc (Private, Lun, Packet, Trd, &PacketBufferMap);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  
    CmdDescBase = (UINT8 *)(UINTN)(LShiftU64 ((UINT64)Trd->UcdBaU, 32) | LShiftU64 ((UINT64)Trd->UcdBa, 7));
    CmdDescSize = Trd->PrdtO * sizeof (UINT32) + Trd->PrdtL * sizeof (UTP_TR_PRD);
  
    //
    // Start to execute the transfer request.
    //
    UfsStartExecCmd (Private, Slot);
  
    //
    // Wait for the completion of the transfer request.
    //
    Address = Private->UfsHcBase + UFS_HC_UTRLDBR_OFFSET;
    Status  = UfsWaitMemSet (Address, BIT0 << Slot, 0, Packet->Timeout);
    if (EFI_ERROR (Status)) {
      goto Exit;
    }
  
    //
    // Get sense data if exists
    //
    Response     = (UTP_RESPONSE_UPIU *)(CmdDescBase + Trd->RuO * sizeof (UINT32));
    SenseDataLen = Response->SenseDataLen;
    SwapLittleEndianToBigEndian ((UINT8 *)&SenseDataLen, sizeof (UINT16));
  
    if ((Packet->SenseDataLength != 0) && (Packet->SenseData != NULL)) {
      //
      // Make sure the hardware device does not return more data than expected.
      //
      if (SenseDataLen <= Packet->SenseDataLength) {
        CopyMem (Packet->SenseData, Response->SenseData, SenseDataLen);
        Packet->SenseDataLength = (UINT8)SenseDataLen;
      } else {
        Packet->SenseDataLength = 0;
      }
    }
  
    //
    // Check the transfer request result.
    //
    if (Response->Response != 0) {
      DEBUG ((DEBUG_ERROR, "UfsExecScsiCmds() fails with Target Failure\n"));
      Status = EFI_DEVICE_ERROR;
      goto Exit;
    }
  
    if (Trd->Ocs == 0) {
      if (Packet->DataDirection == UfsDataIn) {
        if ((Response->Flags & BIT5) == BIT5) {
          ResTranCount = Response->ResTranCount;
          SwapLittleEndianToBigEndian ((UINT8 *)&ResTranCount, sizeof (UINT32));
          Packet->InTransferLength -= ResTranCount;
        }
      } else if (Packet->DataDirection == UfsDataOut) {
        if ((Response->Flags & BIT5) == BIT5) {
          ResTranCount = Response->ResTranCount;
          SwapLittleEndianToBigEndian ((UINT8 *)&ResTranCount, sizeof (UINT32));
          Packet->OutTransferLength -= ResTranCount;
        }
      }
    } else {
      Status = EFI_DEVICE_ERROR;
    }
  
  Exit:
    if (PacketBufferMap != NULL) {
      IoMmuUnmap (PacketBufferMap);
    }
  
    UfsStopExecCmd (Private, Slot);
    UfsPeimFreeMem (Private->Pool, CmdDescBase, CmdDescSize);
  
    return Status;
  }
  
  /**
    Sent UIC DME_LINKSTARTUP command to start the link startup procedure.
  
    @param[in] Private          The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
    @param[in] UicOpcode        The opcode of the UIC command.
    @param[in] Arg1             The value for 1st argument of the UIC command.
    @param[in] Arg2             The value for 2nd argument of the UIC command.
    @param[in] Arg3             The value for 3rd argument of the UIC command.
  
    @return EFI_SUCCESS      Successfully execute this UIC command and detect attached UFS device.
    @return EFI_DEVICE_ERROR Fail to execute this UIC command and detect attached UFS device.
    @return EFI_NOT_FOUND    The presence of the UFS device isn't detected.
  
  **/
  EFI_STATUS
  UfsExecUicCommands (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private,
    IN  UINT8                     UicOpcode,
    IN  UINT32                    Arg1,
    IN  UINT32                    Arg2,
    IN  UINT32                    Arg3
    )
  {
    EFI_STATUS  Status;
    UINTN       Address;
    UINT32      Data;
    UINTN       UfsHcBase;
  
    UfsHcBase = Private->UfsHcBase;
    Address   = UfsHcBase + UFS_HC_IS_OFFSET;
    Data      = MmioRead32 (Address);
    if ((Data & UFS_HC_IS_UCCS) == UFS_HC_IS_UCCS) {
      //
      // Clear IS.BIT10 UIC Command Completion Status (UCCS) at first.
      //
      MmioWrite32 (Address, Data);
    }
  
    //
    // When programming UIC command registers, host software shall set the register UICCMD
    // only after all the UIC command argument registers (UICCMDARG1, UICCMDARG2 and UICCMDARG3)
    // are set.
    //
    Address = UfsHcBase + UFS_HC_UCMD_ARG1_OFFSET;
    MmioWrite32 (Address, Arg1);
  
    Address = UfsHcBase + UFS_HC_UCMD_ARG2_OFFSET;
    MmioWrite32 (Address, Arg2);
  
    Address = UfsHcBase + UFS_HC_UCMD_ARG3_OFFSET;
    MmioWrite32 (Address, Arg3);
  
    //
    // Host software shall only set the UICCMD if HCS.UCRDY is set to 1.
    //
    Address = Private->UfsHcBase + UFS_HC_STATUS_OFFSET;
    Status  = UfsWaitMemSet (Address, UFS_HC_HCS_UCRDY, UFS_HC_HCS_UCRDY, UFS_TIMEOUT);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  
    Address = UfsHcBase + UFS_HC_UIC_CMD_OFFSET;
    MmioWrite32 (Address, (UINT32)UicOpcode);
  
    //
    // UFS 2.0 spec section 5.3.1 Offset:0x20 IS.Bit10 UIC Command Completion Status (UCCS)
    // This bit is set to '1' by the host controller upon completion of a UIC command.
    //
    Address = UfsHcBase + UFS_HC_IS_OFFSET;
    Data    = MmioRead32 (Address);
    Status  = UfsWaitMemSet (Address, UFS_HC_IS_UCCS, UFS_HC_IS_UCCS, UFS_TIMEOUT);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  
    if (UicOpcode != UfsUicDmeReset) {
      Address = UfsHcBase + UFS_HC_UCMD_ARG2_OFFSET;
      Data    = MmioRead32 (Address);
      if ((Data & 0xFF) != 0) {
        DEBUG_CODE_BEGIN ();
        DumpUicCmdExecResult (UicOpcode, (UINT8)(Data & 0xFF));
        DEBUG_CODE_END ();
        return EFI_DEVICE_ERROR;
      }
    }
  
    return EFI_SUCCESS;
  }
  
  /**
    Enable the UFS host controller for accessing.
  
    @param[in] Private                 The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
  
    @retval EFI_SUCCESS                The UFS host controller enabling was executed successfully.
    @retval EFI_DEVICE_ERROR           A device error occurred while enabling the UFS host controller.
  
  **/
  EFI_STATUS
  UfsEnableHostController (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private
    )
  {
    EFI_STATUS  Status;
    UINTN       Address;
    UINT32      Data;
  
    //
    // UFS 2.0 spec section 7.1.1 - Host Controller Initialization
    //
    // Reinitialize the UFS host controller if HCE bit of HC register is set.
    //
    Address = Private->UfsHcBase + UFS_HC_ENABLE_OFFSET;
    Data    = MmioRead32 (Address);
    if ((Data & UFS_HC_HCE_EN) == UFS_HC_HCE_EN) {
      //
      // Write a 0 to the HCE register at first to disable the host controller.
      //
      MmioWrite32 (Address, 0);
      //
      // Wait until HCE is read as '0' before continuing.
      //
      Status = UfsWaitMemSet (Address, UFS_HC_HCE_EN, 0, UFS_TIMEOUT);
      if (EFI_ERROR (Status)) {
        return EFI_DEVICE_ERROR;
      }
    }
  
    //
    // Write a 1 to the HCE register to enable the UFS host controller.
    //
    MmioWrite32 (Address, UFS_HC_HCE_EN);
    //
    // Wait until HCE is read as '1' before continuing.
    //
    Status = UfsWaitMemSet (Address, UFS_HC_HCE_EN, UFS_HC_HCE_EN, UFS_TIMEOUT);
    if (EFI_ERROR (Status)) {
      return EFI_DEVICE_ERROR;
    }
  
    return EFI_SUCCESS;
  }
  
  /**
    Detect if a UFS device attached.
  
    @param[in] Private                 The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
  
    @retval EFI_SUCCESS                The UFS device detection was executed successfully.
    @retval EFI_NOT_FOUND              Not found a UFS device attached.
    @retval EFI_DEVICE_ERROR           A device error occurred while detecting the UFS device.
  
  **/
  EFI_STATUS
  UfsDeviceDetection (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private
    )
  {
    UINTN       Retry;
    UINTN       Address;
    UINT32      Data;
    EFI_STATUS  Status;
  
    //
    // Start UFS device detection.
    // Try up to 3 times for establishing data link with device.
    //
    for (Retry = 0; Retry < 3; Retry++) {
      Status = UfsExecUicCommands (Private, UfsUicDmeLinkStartup, 0, 0, 0);
      if (EFI_ERROR (Status)) {
        return EFI_DEVICE_ERROR;
      }
  
      //
      // Check value of HCS.DP and make sure that there is a device attached to the Link
      //
      Address = Private->UfsHcBase + UFS_HC_STATUS_OFFSET;
      Data    = MmioRead32 (Address);
      if ((Data & UFS_HC_HCS_DP) == 0) {
        Address = Private->UfsHcBase + UFS_HC_IS_OFFSET;
        Status  = UfsWaitMemSet (Address, UFS_HC_IS_ULSS, UFS_HC_IS_ULSS, UFS_TIMEOUT);
        if (EFI_ERROR (Status)) {
          return EFI_DEVICE_ERROR;
        }
      } else {
        DEBUG ((DEBUG_INFO, "UfsblockioPei: found a attached UFS device\n"));
        return EFI_SUCCESS;
      }
    }
  
    return EFI_NOT_FOUND;
  }
  
  /**
    Initialize UFS task management request list related h/w context.
  
    @param[in] Private                 The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
  
    @retval EFI_SUCCESS                The UFS task management list was initialzed successfully.
    @retval EFI_DEVICE_ERROR           The initialization fails.
  
  **/
  EFI_STATUS
  UfsInitTaskManagementRequestList (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private
    )
  {
    UINTN                 Address;
    UINT32                Data;
    UINT8                 Nutmrs;
    VOID                  *CmdDescHost;
    EFI_PHYSICAL_ADDRESS  CmdDescPhyAddr;
    VOID                  *CmdDescMapping;
    EFI_STATUS            Status;
  
    //
    // Initial h/w and s/w context for future operations.
    //
    Address               = Private->UfsHcBase + UFS_HC_CAP_OFFSET;
    Data                  = MmioRead32 (Address);
    Private->Capabilities = Data;
  
    //
    // Allocate and initialize UTP Task Management Request List.
    //
    Nutmrs = (UINT8)(RShiftU64 ((Private->Capabilities & UFS_HC_CAP_NUTMRS), 16) + 1);
    Status = IoMmuAllocateBuffer (
               EFI_SIZE_TO_PAGES (Nutmrs * sizeof (UTP_TMRD)),
               &CmdDescHost,
               &CmdDescPhyAddr,
               &CmdDescMapping
               );
    if (EFI_ERROR (Status)) {
      return EFI_DEVICE_ERROR;
    }
  
    ZeroMem (CmdDescHost, EFI_PAGES_TO_SIZE (EFI_SIZE_TO_PAGES (Nutmrs * sizeof (UTP_TMRD))));
  
    //
    // Program the UTP Task Management Request List Base Address and UTP Task Management
    // Request List Base Address with a 64-bit address allocated at step 6.
    //
    Address = Private->UfsHcBase + UFS_HC_UTMRLBA_OFFSET;
    MmioWrite32 (Address, (UINT32)(UINTN)CmdDescPhyAddr);
    Address = Private->UfsHcBase + UFS_HC_UTMRLBAU_OFFSET;
    MmioWrite32 (Address, (UINT32)RShiftU64 ((UINT64)CmdDescPhyAddr, 32));
    Private->UtpTmrlBase = (VOID *)(UINTN)CmdDescHost;
    Private->Nutmrs      = Nutmrs;
    Private->TmrlMapping = CmdDescMapping;
  
    //
    // Enable the UTP Task Management Request List by setting the UTP Task Management
    // Request List RunStop Register (UTMRLRSR) to '1'.
    //
    Address = Private->UfsHcBase + UFS_HC_UTMRLRSR_OFFSET;
    MmioWrite32 (Address, UFS_HC_UTMRLRSR);
  
    return EFI_SUCCESS;
  }
  
  /**
    Initialize UFS transfer request list related h/w context.
  
    @param[in] Private                 The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
  
    @retval EFI_SUCCESS                The UFS transfer list was initialzed successfully.
    @retval EFI_DEVICE_ERROR           The initialization fails.
  
  **/
  EFI_STATUS
  UfsInitTransferRequestList (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private
    )
  {
    UINTN                 Address;
    UINT32                Data;
    UINT8                 Nutrs;
    VOID                  *CmdDescHost;
    EFI_PHYSICAL_ADDRESS  CmdDescPhyAddr;
    VOID                  *CmdDescMapping;
    EFI_STATUS            Status;
  
    //
    // Initial h/w and s/w context for future operations.
    //
    Address               = Private->UfsHcBase + UFS_HC_CAP_OFFSET;
    Data                  = MmioRead32 (Address);
    Private->Capabilities = Data;
  
    //
    // Allocate and initialize UTP Transfer Request List.
    //
    Nutrs  = (UINT8)((Private->Capabilities & UFS_HC_CAP_NUTRS) + 1);
    Status = IoMmuAllocateBuffer (
               EFI_SIZE_TO_PAGES (Nutrs * sizeof (UTP_TRD)),
               &CmdDescHost,
               &CmdDescPhyAddr,
               &CmdDescMapping
               );
    if (EFI_ERROR (Status)) {
      return EFI_DEVICE_ERROR;
    }
  
    ZeroMem (CmdDescHost, EFI_PAGES_TO_SIZE (EFI_SIZE_TO_PAGES (Nutrs * sizeof (UTP_TRD))));
  
    //
    // Program the UTP Transfer Request List Base Address and UTP Transfer Request List
    // Base Address with a 64-bit address allocated at step 8.
    //
    Address = Private->UfsHcBase + UFS_HC_UTRLBA_OFFSET;
    MmioWrite32 (Address, (UINT32)(UINTN)CmdDescPhyAddr);
    Address = Private->UfsHcBase + UFS_HC_UTRLBAU_OFFSET;
    MmioWrite32 (Address, (UINT32)RShiftU64 ((UINT64)CmdDescPhyAddr, 32));
    Private->UtpTrlBase = (VOID *)(UINTN)CmdDescHost;
    Private->Nutrs      = Nutrs;
    Private->TrlMapping = CmdDescMapping;
  
    //
    // Enable the UTP Transfer Request List by setting the UTP Transfer Request List
    // RunStop Register (UTRLRSR) to '1'.
    //
    Address = Private->UfsHcBase + UFS_HC_UTRLRSR_OFFSET;
    MmioWrite32 (Address, UFS_HC_UTRLRSR);
  
    return EFI_SUCCESS;
  }
  
  /**
    Initialize the UFS host controller.
  
    @param[in] Private                 The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
  
    @retval EFI_SUCCESS                The Ufs Host Controller is initialized successfully.
    @retval Others                     A device error occurred while initializing the controller.
  
  **/
  EFI_STATUS
  UfsControllerInit (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private
    )
  {
    EFI_STATUS  Status;
  
    Status = UfsEnableHostController (Private);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "UfsDevicePei: Enable Host Controller Fails, Status = %r\n", Status));
      return Status;
    }
  
    Status = UfsDeviceDetection (Private);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "UfsDevicePei: Device Detection Fails, Status = %r\n", Status));
      return Status;
    }
  
    Status = UfsInitTaskManagementRequestList (Private);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "UfsDevicePei: Task management list initialization Fails, Status = %r\n", Status));
      return Status;
    }
  
    Status = UfsInitTransferRequestList (Private);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "UfsDevicePei: Transfer list initialization Fails, Status = %r\n", Status));
  
      if (Private->TmrlMapping != NULL) {
        IoMmuFreeBuffer (
          EFI_SIZE_TO_PAGES (Private->Nutmrs * sizeof (UTP_TMRD)),
          Private->UtpTmrlBase,
          Private->TmrlMapping
          );
        Private->TmrlMapping = NULL;
      }
  
      return Status;
    }
  
    DEBUG ((DEBUG_INFO, "UfsDevicePei Finished\n"));
    return EFI_SUCCESS;
  }
  
  /**
    Stop the UFS host controller.
  
    @param[in] Private                 The pointer to the UFS_PEIM_HC_PRIVATE_DATA data structure.
  
    @retval EFI_SUCCESS                The Ufs Host Controller is stopped successfully.
    @retval Others                     A device error occurred while stopping the controller.
  
  **/
  EFI_STATUS
  UfsControllerStop (
    IN  UFS_PEIM_HC_PRIVATE_DATA  *Private
    )
  {
    EFI_STATUS  Status;
    UINTN       Address;
    UINT32      Data;
  
    //
    // Enable the UTP Task Management Request List by setting the UTP Task Management
    // Request List RunStop Register (UTMRLRSR) to '1'.
    //
    Address = Private->UfsHcBase + UFS_HC_UTMRLRSR_OFFSET;
    MmioWrite32 (Address, 0);
  
    //
    // Enable the UTP Transfer Request List by setting the UTP Transfer Request List
    // RunStop Register (UTRLRSR) to '1'.
    //
    Address = Private->UfsHcBase + UFS_HC_UTRLRSR_OFFSET;
    MmioWrite32 (Address, 0);
  
    //
    // Write a 0 to the HCE register in order to disable the host controller.
    //
    Address = Private->UfsHcBase + UFS_HC_ENABLE_OFFSET;
    Data    = MmioRead32 (Address);
    ASSERT ((Data & UFS_HC_HCE_EN) == UFS_HC_HCE_EN);
    MmioWrite32 (Address, 0);
  
    //
    // Wait until HCE is read as '0' before continuing.
    //
    Status = UfsWaitMemSet (Address, UFS_HC_HCE_EN, 0, UFS_TIMEOUT);
    if (EFI_ERROR (Status)) {
      return EFI_DEVICE_ERROR;
    }
  
    DEBUG ((DEBUG_INFO, "UfsDevicePei: Stop the UFS Host Controller\n"));
  
    return EFI_SUCCESS;
  }
  