[PATCH 2/2] MdeModulePkg/Ufs: Wait fDeviceInit be cleared by devices during init

[Hao Wu <hao.a.wu@intel.com>]

In the origin codes, the host sets the fDeviceInit flag to initiate device
initialization, but does not check whether the device resets this flag
to indicate the device initialization is completed.

Details can be referred at UFS 2.0 Spec Section 14.2 - Flags.

Cc: Feng Tian <feng.tian@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Mateusz Albecki <mateusz.albecki@intel.com>
Signed-off-by: Hao Wu <hao.a.wu@intel.com>
---
 MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.c    |  69 +++-
 MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.h    |  53 ++-
 MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThruHci.c | 401 ++++++++------------
 3 files changed, 267 insertions(+), 256 deletions(-)

diff --git a/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.c b/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.c
index 63b90de991..68a44367b5 100644
--- a/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.c
+++ b/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.c
@@ -730,6 +730,48 @@ UfsPassThruDriverBindingSupported (
 }
 
 /**
+  Finishes device initialization by setting fDeviceInit flag and waiting untill device responds by
+  clearing it.
+
+  @param[in] Private  Pointer to the UFS_PASS_THRU_PRIVATE_DATA.
+
+  @retval EFI_SUCCESS  The operation succeeds.
+  @retval Others       The operation fails.
+
+**/
+EFI_STATUS
+UfsFinishDeviceInitialization (
+  IN UFS_PASS_THRU_PRIVATE_DATA  *Private
+  )
+{
+  EFI_STATUS  Status;
+  UINT8  DeviceInitStatus;
+  UINT8  Timeout;
+
+  DeviceInitStatus = 0xFF;
+
+  //
+  // The host enables the device initialization completion by setting fDeviceInit flag.
+  //
+  Status = UfsSetFlag (Private, UfsFlagDevInit);
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  Timeout = 5;
+  do {
+    Status = UfsReadFlag (Private, UfsFlagDevInit, &DeviceInitStatus);
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
+    MicroSecondDelay (1);
+    Timeout--;
+  } while (DeviceInitStatus != 0 && Timeout != 0);
+
+  return EFI_SUCCESS;
+}
+
+/**
   Starts a device controller or a bus controller.
 
   The Start() function is designed to be invoked from the EFI boot service ConnectController().
@@ -777,7 +819,7 @@ UfsPassThruDriverBindingStart (
   UFS_PASS_THRU_PRIVATE_DATA            *Private;
   UINTN                                 UfsHcBase;
   UINT32                                Index;
-  UFS_CONFIG_DESC                       Config;
+  UFS_UNIT_DESC                         UnitDescriptor;
 
   Status    = EFI_SUCCESS;
   UfsHc     = NULL;
@@ -844,21 +886,9 @@ UfsPassThruDriverBindingStart (
     goto Error;
   }
 
-  //
-  // The host enables the device initialization completion by setting fDeviceInit flag.
-  //
-  Status = UfsSetFlag (Private, UfsFlagDevInit);
+  Status = UfsFinishDeviceInitialization (Private);
   if (EFI_ERROR (Status)) {
-    DEBUG ((DEBUG_ERROR, "Ufs Set fDeviceInit Flag Error, Status = %r\n", Status));
-    goto Error;
-  }
-
-  //
-  // Get Ufs Device's Lun Info by reading Configuration Descriptor.
-  //
-  Status = UfsRwDeviceDesc (Private, TRUE, UfsConfigDesc, 0, 0, &Config, sizeof (UFS_CONFIG_DESC));
-  if (EFI_ERROR (Status)) {
-    DEBUG ((DEBUG_ERROR, "Ufs Get Configuration Descriptor Error, Status = %r\n", Status));
+    DEBUG ((DEBUG_ERROR, "Device failed to finish initialization, Status = %r\n", Status));
     goto Error;
   }
 
@@ -867,9 +897,14 @@ UfsPassThruDriverBindingStart (
   // TODO: Parse device descriptor to decide if exposing RPMB LUN to upper layer for authentication access.
   //
   for (Index = 0; Index < 8; Index++) {
-    if (Config.UnitDescConfParams[Index].LunEn != 0) {
+    Status = UfsRwDeviceDesc (Private, TRUE, UfsUnitDesc, (UINT8) Index, 0, &UnitDescriptor, sizeof (UFS_UNIT_DESC));
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "Failed to read unit descriptor, index = %X, status = %r\n", Index, Status));
+      continue;
+    }
+    if (UnitDescriptor.LunEn == 0x1) {
+      DEBUG ((DEBUG_INFO, "UFS LUN %X is enabled\n", Index));
       Private->Luns.BitMask |= (BIT0 << Index);
-      DEBUG ((DEBUG_INFO, "Ufs Lun %d is enabled\n", Index));
     }
   }
 
diff --git a/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.h b/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.h
index 7fc82bae7a..bdc64f7793 100644
--- a/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.h
+++ b/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThru.h
@@ -1,6 +1,6 @@
 /** @file
 
-  Copyright (c) 2014 - 2015, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2014 - 2017, Intel Corporation. All rights reserved.<BR>
   This program and the accompanying materials
   are licensed and made available under the terms and conditions of the BSD License
   which accompanies this distribution.  The full text of the license may be found at
@@ -122,16 +122,13 @@ typedef struct {
 
 typedef struct _UFS_DEVICE_MANAGEMENT_REQUEST_PACKET {
   UINT64           Timeout;
-  VOID             *InDataBuffer;
-  VOID             *OutDataBuffer;
+  VOID             *DataBuffer;
   UINT8            Opcode;
   UINT8            DescId;
   UINT8            Index;
   UINT8            Selector;
-  UINT32           InTransferLength;
-  UINT32           OutTransferLength;
+  UINT32           TransferLength;
   UINT8            DataDirection;
-  UINT8            Ocs;
 } UFS_DEVICE_MANAGEMENT_REQUEST_PACKET;
 
 //
@@ -717,6 +714,25 @@ UfsSetFlag (
   );
 
 /**
+  Read specified flag from a UFS device.
+
+  @param[in]  Private           The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.
+  @param[in]  FlagId            The ID of flag to be read.
+  @param[out] Value             The flag's value.
+
+  @retval EFI_SUCCESS           The flag was read successfully.
+  @retval EFI_DEVICE_ERROR      A device error occurred while attempting to read the flag.
+  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of reading the flag.
+
+**/
+EFI_STATUS
+UfsReadFlag (
+  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,
+  IN     UINT8                        FlagId,
+     OUT UINT8                        *Value
+  );
+
+/**
   Read or write specified device descriptor of a UFS device.
 
   @param[in]      Private       The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.
@@ -744,6 +760,31 @@ UfsRwDeviceDesc (
   );
 
 /**
+  Read or write specified attribute of a UFS device.
+
+  @param[in]      Private       The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.
+  @param[in]      Read          The boolean variable to show r/w direction.
+  @param[in]      AttrId        The ID of Attribute.
+  @param[in]      Index         The Index of Attribute.
+  @param[in]      Selector      The Selector of Attribute.
+  @param[in, out] Attributes    The value of Attribute to be read or written.
+
+  @retval EFI_SUCCESS           The Attribute was read/written successfully.
+  @retval EFI_DEVICE_ERROR      A device error occurred while attempting to r/w the Attribute.
+  @retval EFI_TIMEOUT           A timeout occurred while waiting for the completion of r/w the Attribute.
+
+**/
+EFI_STATUS
+UfsRwAttributes (
+  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,
+  IN     BOOLEAN                      Read,
+  IN     UINT8                        AttrId,
+  IN     UINT8                        Index,
+  IN     UINT8                        Selector,
+  IN OUT UINT32                       *Attributes
+  );
+
+/**
   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.
 
diff --git a/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThruHci.c b/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThruHci.c
index 822554cebb..a19bdcc3cc 100644
--- a/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThruHci.c
+++ b/MdeModulePkg/Bus/Ufs/UfsPassThruDxe/UfsPassThruHci.c
@@ -602,32 +602,13 @@ UfsCreateDMCommandDesc (
   }
 
   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 != UtpQueryFuncOpcodeSetFlag) && (Opcode != UtpQueryFuncOpcodeClrFlag) && (Opcode != UtpQueryFuncOpcodeTogFlag))
-    && ((DataSize == 0) || (Data == NULL))) {
-    return EFI_INVALID_PARAMETER;
-  }
-
-  if (((Opcode == UtpQueryFuncOpcodeSetFlag) || (Opcode == UtpQueryFuncOpcodeClrFlag) || (Opcode == UtpQueryFuncOpcodeTogFlag))
-    && ((DataSize != 0) || (Data != NULL))) {
-    return EFI_INVALID_PARAMETER;
-  }
-
-  if ((Opcode == UtpQueryFuncOpcodeWrAttr) && (DataSize != sizeof (UINT32))) {
-    return EFI_INVALID_PARAMETER;
-  }
+  DataSize      = Packet->TransferLength;
+  Data          = Packet->DataBuffer;
 
   if ((Opcode == UtpQueryFuncOpcodeWrDesc) || (Opcode == UtpQueryFuncOpcodeRdDesc)) {
+    if (DataSize == 0 || Data == NULL) {
+      return EFI_INVALID_PARAMETER;
+    }
     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));
@@ -868,60 +849,84 @@ UfsStopExecCmd (
 }
 
 /**
-  Read or write specified device descriptor of a UFS device.
+  Extracts return data from query response upiu.
 
-  @param[in]      Private       The pointer to the UFS_PASS_THRU_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.
+  @param[in] Packet     Pointer to the UFS_DEVICE_MANAGEMENT_REQUEST_PACKET.
+  @param[in] QueryResp  Pointer to the query response.
 
-  @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.
+  @retval EFI_INVALID_PARAMETER Packet or QueryResp are empty or opcode is invalid.
+  @retval EFI_SUCCESS           Data extracted.
 
 **/
 EFI_STATUS
-UfsRwDeviceDesc (
-  IN     UFS_PASS_THRU_PRIVATE_DATA   *Private,
-  IN     BOOLEAN                      Read,
-  IN     UINT8                        DescId,
-  IN     UINT8                        Index,
-  IN     UINT8                        Selector,
-  IN OUT VOID                         *Descriptor,
-  IN     UINT32                       DescSize
+UfsGetReturnDataFromQueryResponse (
+  IN UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet,
+  IN UTP_QUERY_RESP_UPIU                   *QueryResp
   )
 {
-  EFI_STATUS                           Status;
-  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET Packet;
-  UINT8                                Slot;
-  UTP_TRD                              *Trd;
-  UTP_QUERY_RESP_UPIU                  *QueryResp;
-  UINT32                               CmdDescSize;
-  UINT16                               ReturnDataSize;
-  VOID                                 *CmdDescHost;
-  VOID                                 *CmdDescMapping;
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;
+  UINT16  ReturnDataSize;
+  UINT32  ReturnData;
 
-  ZeroMem (&Packet, sizeof (UFS_DEVICE_MANAGEMENT_REQUEST_PACKET));
+  if (Packet == NULL || QueryResp == NULL) {
+    return EFI_INVALID_PARAMETER;
+  }
 
-  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;
+  switch (Packet->Opcode) {
+    case UtpQueryFuncOpcodeRdDesc:
+      ReturnDataSize = QueryResp->Tsf.Length;
+      SwapLittleEndianToBigEndian ((UINT8*)&ReturnDataSize, sizeof (UINT16));
+      CopyMem (Packet->DataBuffer, (QueryResp + 1), ReturnDataSize);
+      Packet->TransferLength = ReturnDataSize;
+      break;
+    case UtpQueryFuncOpcodeWrDesc:
+      ReturnDataSize = QueryResp->Tsf.Length;
+      SwapLittleEndianToBigEndian ((UINT8*)&ReturnDataSize, sizeof (UINT16));
+      Packet->TransferLength = ReturnDataSize;
+      break;
+    case UtpQueryFuncOpcodeRdFlag:
+    case UtpQueryFuncOpcodeSetFlag:
+    case UtpQueryFuncOpcodeClrFlag:
+    case UtpQueryFuncOpcodeTogFlag:
+      CopyMem (Packet->DataBuffer, &QueryResp->Tsf.Value, sizeof (UINT8));
+      break;
+    case UtpQueryFuncOpcodeRdAttr:
+    case UtpQueryFuncOpcodeWrAttr:
+      ReturnData = QueryResp->Tsf.Value;
+      SwapLittleEndianToBigEndian ((UINT8*) &ReturnData, sizeof (UINT32));
+      CopyMem (Packet->DataBuffer, &ReturnData, sizeof (UINT32));
+      break;
+    default:
+      return EFI_INVALID_PARAMETER;
   }
-  Packet.DescId              = DescId;
-  Packet.Index               = Index;
-  Packet.Selector            = Selector;
-  Packet.Timeout             = UFS_TIMEOUT;
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Creates Transfer Request descriptor and sends Query Request to the device.
+
+  @param[in] Private Pointer to the UFS_PASS_THRU_PRIVATE_DATA.
+  @param[in] Packet  Pointer to the UFS_DEVICE_MANAGEMENT_REQUEST_PACKET.
+
+  @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
+UfsSendDmRequestRetry (
+  IN UFS_PASS_THRU_PRIVATE_DATA            *Private,
+  IN UFS_DEVICE_MANAGEMENT_REQUEST_PACKET  *Packet
+  )
+{
+  UINT8                               Slot;
+  UTP_TRD                             *Trd;
+  VOID                                *CmdDescHost;
+  VOID                                *CmdDescMapping;
+  UINT32                              CmdDescSize;
+  EDKII_UFS_HOST_CONTROLLER_PROTOCOL  *UfsHc;
+  UTP_QUERY_RESP_UPIU                 *QueryResp;
+  EFI_STATUS                          Status;
 
   //
   // Find out which slot of transfer request list is available.
@@ -935,14 +940,12 @@ UfsRwDeviceDesc (
   //
   // Fill transfer request descriptor to this slot.
   //
-  Status = UfsCreateDMCommandDesc (Private, &Packet, Trd, &CmdDescHost, &CmdDescMapping);
+  Status = UfsCreateDMCommandDesc (Private, Packet, Trd, &CmdDescHost, &CmdDescMapping);
   if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Failed to create DM command descriptor\n"));
     return Status;
   }
 
-  //
-  // Check the transfer request result.
-  //
   UfsHc       = Private->UfsHostController;
   QueryResp   = (UTP_QUERY_RESP_UPIU*)((UINT8*)CmdDescHost + Trd->RuO * sizeof (UINT32));
   ASSERT (QueryResp != NULL);
@@ -955,30 +958,23 @@ UfsRwDeviceDesc (
 
   //
   // Wait for the completion of the transfer request.
-  //  
-  Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0 << Slot, 0, Packet.Timeout);
+  //
+  Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0, 0, Packet->Timeout);
   if (EFI_ERROR (Status)) {
     goto Exit;
   }
 
-  if (QueryResp->QueryResp != 0) {
+  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;
   }
 
-  if (Trd->Ocs == 0) {
-    ReturnDataSize = QueryResp->Tsf.Length;
-    SwapLittleEndianToBigEndian ((UINT8*)&ReturnDataSize, sizeof (UINT16));
-
-    if (Read) {
-      CopyMem (Packet.InDataBuffer, (QueryResp + 1), ReturnDataSize);
-      Packet.InTransferLength = ReturnDataSize;
-    } else {
-      Packet.OutTransferLength = ReturnDataSize;
-    }
-  } else {
-    Status = EFI_DEVICE_ERROR;
+  Status = UfsGetReturnDataFromQueryResponse (Packet, QueryResp);
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Failed to get return data from query response\n"));
+    goto Exit;
   }
 
 Exit:
@@ -997,6 +993,87 @@ Exit:
 }
 
 /**
+  Sends Query Request to the device. Query is sent until device responds correctly or counter runs out.
+
+  @param[in] Private Pointer to the UFS_PASS_THRU_PRIVATE_DATA.
+  @param[in] Packet  Pointer to the UFS_DEVICE_MANAGEMENT_PACKET.
+
+  @retval EFI_SUCCESS           The device responded correctly to the Query request.
+  @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_PASS_THRU_PRIVATE_DATA            *Private,
+  IN 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_PASS_THRU_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_PASS_THRU_PRIVATE_DATA   *Private,
+  IN     BOOLEAN                      Read,
+  IN     UINT8                        DescId,
+  IN     UINT8                        Index,
+  IN     UINT8                        Selector,
+  IN OUT VOID                         *Descriptor,
+  IN     UINT32                       DescSize
+  )
+{
+  UFS_DEVICE_MANAGEMENT_REQUEST_PACKET Packet;
+
+  ZeroMem (&Packet, sizeof (UFS_DEVICE_MANAGEMENT_REQUEST_PACKET));
+
+  if (Read) {
+    Packet.DataDirection     = UfsDataIn;
+    Packet.Opcode            = UtpQueryFuncOpcodeRdDesc;
+  } else {
+    Packet.DataDirection     = UfsDataOut;
+    Packet.Opcode            = UtpQueryFuncOpcodeWrDesc;
+  }
+  Packet.DataBuffer          = Descriptor;
+  Packet.TransferLength      = DescSize;
+  Packet.DescId              = DescId;
+  Packet.Index               = Index;
+  Packet.Selector            = Selector;
+  Packet.Timeout             = UFS_TIMEOUT;
+
+  return UfsSendDmRequest (Private, &Packet);
+}
+
+/**
   Read or write specified attribute of a UFS device.
 
   @param[in]      Private       The pointer to the UFS_PASS_THRU_PRIVATE_DATA data structure.
@@ -1021,16 +1098,7 @@ UfsRwAttributes (
   IN OUT UINT32                       *Attributes
   )
 {
-  EFI_STATUS                           Status;
   UFS_DEVICE_MANAGEMENT_REQUEST_PACKET Packet;
-  UINT8                                Slot;
-  UTP_TRD                              *Trd;
-  UTP_QUERY_RESP_UPIU                  *QueryResp;
-  UINT32                               CmdDescSize;
-  UINT32                               ReturnData;
-  VOID                                 *CmdDescHost;
-  VOID                                 *CmdDescMapping;
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;
 
   ZeroMem (&Packet, sizeof (UFS_DEVICE_MANAGEMENT_REQUEST_PACKET));
 
@@ -1041,77 +1109,13 @@ UfsRwAttributes (
     Packet.DataDirection     = UfsDataOut;
     Packet.Opcode            = UtpQueryFuncOpcodeWrAttr;
   }
+  Packet.DataBuffer          = Attributes;
   Packet.DescId              = AttrId;
   Packet.Index               = Index;
   Packet.Selector            = Selector;
   Packet.Timeout             = UFS_TIMEOUT;
 
-  //
-  // 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, &CmdDescHost, &CmdDescMapping);
-  if (EFI_ERROR (Status)) {
-    return Status;
-  }
-
-  //
-  // Check the transfer request result.
-  //
-  UfsHc       = Private->UfsHostController;
-  QueryResp   = (UTP_QUERY_RESP_UPIU*)((UINT8*)CmdDescHost + Trd->RuO * sizeof (UINT32));
-  ASSERT (QueryResp != NULL);
-  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.
-  //  
-  Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0 << Slot, 0, Packet.Timeout);
-  if (EFI_ERROR (Status)) {
-    goto Exit;
-  }
-
-  if (QueryResp->QueryResp != 0) {
-    DumpQueryResponseResult (QueryResp->QueryResp);
-    Status = EFI_DEVICE_ERROR;
-    goto Exit;
-  }
-
-  if (Trd->Ocs == 0) {
-    ReturnData = QueryResp->Tsf.Value;
-    SwapLittleEndianToBigEndian ((UINT8*)&ReturnData, sizeof (UINT32));
-    *Attributes = ReturnData;
-  } else {
-    Status = EFI_DEVICE_ERROR;
-  }
-
-Exit:
-  UfsHc->Flush (UfsHc);
-
-  UfsStopExecCmd (Private, Slot);
-
-  if (CmdDescMapping != NULL) {
-    UfsHc->Unmap (UfsHc, CmdDescMapping);
-  }
-
-  if (CmdDescHost != NULL) {
-    UfsHc->FreeBuffer (UfsHc, EFI_SIZE_TO_PAGES (CmdDescSize), CmdDescHost);
-  }
-
-  return Status;
+  return UfsSendDmRequest (Private, &Packet);
 }
 
 /**
@@ -1135,15 +1139,7 @@ UfsRwFlags (
   IN OUT UINT8                        *Value
   )
 {
-  EFI_STATUS                           Status;
   UFS_DEVICE_MANAGEMENT_REQUEST_PACKET Packet;
-  UINT8                                Slot;
-  UTP_TRD                              *Trd;
-  UTP_QUERY_RESP_UPIU                  *QueryResp;
-  UINT32                               CmdDescSize;
-  VOID                                 *CmdDescHost;
-  VOID                                 *CmdDescMapping;
-  EDKII_UFS_HOST_CONTROLLER_PROTOCOL   *UfsHc;
 
   if (Value == NULL) {
     return EFI_INVALID_PARAMETER;
@@ -1165,74 +1161,13 @@ UfsRwFlags (
       return EFI_INVALID_PARAMETER;
     }
   }
+  Packet.DataBuffer          = Value;
   Packet.DescId              = FlagId;
   Packet.Index               = 0;
   Packet.Selector            = 0;
   Packet.Timeout             = UFS_TIMEOUT;
 
-  //
-  // Find out which slot of transfer request list is available.
-  //
-  Status = UfsFindAvailableSlotInTrl (Private, &Slot);
-  if (EFI_ERROR (Status)) {
-    return Status;
-  }
-
-  //
-  // Fill transfer request descriptor to this slot.
-  //
-  Trd    = ((UTP_TRD*)Private->UtpTrlBase) + Slot;
-  Status = UfsCreateDMCommandDesc (Private, &Packet, Trd, &CmdDescHost, &CmdDescMapping);
-  if (EFI_ERROR (Status)) {
-    return Status;
-  }
-
-  //
-  // Check the transfer request result.
-  //
-  UfsHc       = Private->UfsHostController;
-  QueryResp   = (UTP_QUERY_RESP_UPIU*)((UINT8*)CmdDescHost + Trd->RuO * sizeof (UINT32));
-  ASSERT (QueryResp != NULL);
-  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.
-  //  
-  Status = UfsWaitMemSet (Private, UFS_HC_UTRLDBR_OFFSET, BIT0 << Slot, 0, Packet.Timeout);
-  if (EFI_ERROR (Status)) {
-    goto Exit;
-  }
-
-  if (QueryResp->QueryResp != 0) {
-    DumpQueryResponseResult (QueryResp->QueryResp);
-    Status = EFI_DEVICE_ERROR;
-    goto Exit;
-  }
-
-  if (Trd->Ocs == 0) {
-    *Value = (UINT8)QueryResp->Tsf.Value;
-  } else {
-    Status = EFI_DEVICE_ERROR;
-  }
-
-Exit:
-  UfsHc->Flush (UfsHc);
-
-  UfsStopExecCmd (Private, Slot);
-
-  if (CmdDescMapping != NULL) {
-    UfsHc->Unmap (UfsHc, CmdDescMapping);
-  }
-  if (CmdDescHost != NULL) {
-    UfsHc->FreeBuffer (UfsHc, EFI_SIZE_TO_PAGES (CmdDescSize), CmdDescHost);
-  }
-
-  return Status;
+  return UfsSendDmRequest (Private, &Packet);
 }
 
 /**
-- 
2.12.0.windows.1