MdeModulePkg/XhciPei: Fix Aligned Page Allocation

Add support for allocating aligned pages at an alignment higher
than 4K. The new function allocated memory taking into account
the padding required and then frees up unused pages before mapping
it.

Signed-off-by: Ashish Singhal <ashishsingha@nvidia.com>
Reviewed-by: Hao A Wu <hao.a.wu@intel.com>

MdeModulePkg/Bus/Pci/XhciPei/DmaMem.c

@@ -224,6 +224,134 @@ IoMmuFreeBuffer (
   return Status;
 }
 
+/**
+  Allocates aligned pages that are suitable for an OperationBusMasterCommonBuffer or
+  OperationBusMasterCommonBuffer64 mapping.
+
+  @param  Pages                 The number of pages to allocate.
+  @param  Alignment             The requested alignment of the allocation.  Must be a power of two.
+  @param  HostAddress           A pointer to store the base system memory address of the
+                                allocated range.
+  @param  DeviceAddress         The resulting map address for the bus master PCI controller to use to
+                                access the hosts HostAddress.
+  @param  Mapping               A resulting value to pass to Unmap().
+
+  @retval EFI_SUCCESS           The requested memory pages were allocated.
+  @retval EFI_UNSUPPORTED       Attributes is unsupported. The only legal attribute bits are
+                                MEMORY_WRITE_COMBINE and MEMORY_CACHED.
+  @retval EFI_INVALID_PARAMETER One or more parameters are invalid.
+  @retval EFI_OUT_OF_RESOURCES  The memory pages could not be allocated.
+
+**/
+EFI_STATUS
+IoMmuAllocateAlignedBuffer (
+  IN UINTN                  Pages,
+  IN UINTN                  Alignment,
+  OUT VOID                  **HostAddress,
+  OUT EFI_PHYSICAL_ADDRESS  *DeviceAddress,
+  OUT VOID                  **Mapping
+  )
+{
+  EFI_STATUS            Status;
+  VOID                  *Memory;
+  UINTN                 AlignedMemory;
+  UINTN                 AlignmentMask;
+  UINTN                 UnalignedPages;
+  UINTN                 RealPages;
+  UINTN                 NumberOfBytes;
+  EFI_PHYSICAL_ADDRESS  HostPhyAddress;
+
+  *HostAddress   = NULL;
+  *DeviceAddress = 0;
+  AlignmentMask  = Alignment - 1;
+
+  //
+  // Calculate the total number of pages since alignment is larger than page size.
+  //
+  RealPages = Pages + EFI_SIZE_TO_PAGES (Alignment);
+
+  //
+  // Make sure that Pages plus EFI_SIZE_TO_PAGES (Alignment) does not overflow.
+  //
+  ASSERT (RealPages > Pages);
+
+  if (mIoMmu != NULL) {
+    Status = mIoMmu->AllocateBuffer (
+                       mIoMmu,
+                       EfiBootServicesData,
+                       RealPages,
+                       HostAddress,
+                       0
+                       );
+    if (EFI_ERROR (Status)) {
+      return EFI_OUT_OF_RESOURCES;
+    }
+    Memory         = *HostAddress;
+    AlignedMemory  = ((UINTN) Memory + AlignmentMask) & ~AlignmentMask;
+    UnalignedPages = EFI_SIZE_TO_PAGES (AlignedMemory - (UINTN) Memory);
+    if (UnalignedPages > 0) {
+      //
+      // Free first unaligned page(s).
+      //
+      Status = mIoMmu->FreeBuffer (
+                         mIoMmu,
+                         UnalignedPages,
+                         Memory);
+      if (EFI_ERROR (Status)) {
+        return Status;
+      }
+    }
+    Memory         = (VOID *)(UINTN)(AlignedMemory + EFI_PAGES_TO_SIZE (Pages));
+    UnalignedPages = RealPages - Pages - UnalignedPages;
+    if (UnalignedPages > 0) {
+      //
+      // Free last unaligned page(s).
+      //
+      Status = mIoMmu->FreeBuffer (
+                         mIoMmu,
+                         UnalignedPages,
+                         Memory);
+      if (EFI_ERROR (Status)) {
+        return Status;
+      }
+    }
+    *HostAddress  = (VOID *) AlignedMemory;
+    NumberOfBytes = EFI_PAGES_TO_SIZE(Pages);
+    Status = mIoMmu->Map (
+                       mIoMmu,
+                       EdkiiIoMmuOperationBusMasterCommonBuffer,
+                       *HostAddress,
+                       &NumberOfBytes,
+                       DeviceAddress,
+                       Mapping
+                       );
+    if (EFI_ERROR (Status)) {
+      return EFI_OUT_OF_RESOURCES;
+    }
+    Status = mIoMmu->SetAttribute (
+                       mIoMmu,
+                       *Mapping,
+                       EDKII_IOMMU_ACCESS_READ | EDKII_IOMMU_ACCESS_WRITE
+                       );
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
+  } else {
+    Status = PeiServicesAllocatePages (
+               EfiBootServicesData,
+               RealPages,
+               &HostPhyAddress
+               );
+    if (EFI_ERROR (Status)) {
+      return EFI_OUT_OF_RESOURCES;
+    }
+    *HostAddress = (VOID *)(((UINTN) HostPhyAddress + AlignmentMask) & ~AlignmentMask);
+    *DeviceAddress = ((UINTN) HostPhyAddress + AlignmentMask) & ~AlignmentMask;
+    *Mapping = NULL;
+  }
+  return Status;
+}
+
 /**
   Initialize IOMMU.
 **/

MdeModulePkg/Bus/Pci/XhciPei/UsbHcMem.c

@@ -562,11 +562,7 @@ UsbHcAllocateAlignedPages (
 {
   EFI_STATUS            Status;
   VOID                  *Memory;
-  UINTN                 AlignedMemory;
-  UINTN                 AlignmentMask;
   EFI_PHYSICAL_ADDRESS  DeviceMemory;
-  UINTN                 AlignedDeviceMemory;
-  UINTN                 RealPages;
 
   //
   // Alignment must be a power of two or zero.
@@ -582,18 +578,9 @@ UsbHcAllocateAlignedPages (
   }
 
   if (Alignment > EFI_PAGE_SIZE) {
-    //
+    Status = IoMmuAllocateAlignedBuffer (
-    // Calculate the total number of pages since alignment is larger than page size.
-    //
-    AlignmentMask  = Alignment - 1;
-    RealPages      = Pages + EFI_SIZE_TO_PAGES (Alignment);
-    //
-    // Make sure that Pages plus EFI_SIZE_TO_PAGES (Alignment) does not overflow.
-    //
-    ASSERT (RealPages > Pages);
-
-    Status = IoMmuAllocateBuffer (
                Pages,
+               Alignment,
                &Memory,
                &DeviceMemory,
                Mapping
@@ -601,8 +588,6 @@ UsbHcAllocateAlignedPages (
     if (EFI_ERROR (Status)) {
       return EFI_OUT_OF_RESOURCES;
     }
-    AlignedMemory = ((UINTN) Memory + AlignmentMask) & ~AlignmentMask;
-    AlignedDeviceMemory = ((UINTN) DeviceMemory + AlignmentMask) & ~AlignmentMask;
   } else {
     //
     // Do not over-allocate pages in this case.
@@ -616,12 +601,10 @@ UsbHcAllocateAlignedPages (
     if (EFI_ERROR (Status)) {
       return EFI_OUT_OF_RESOURCES;
     }
-    AlignedMemory = (UINTN) Memory;
-    AlignedDeviceMemory = (UINTN) DeviceMemory;
   }
 
-  *HostAddress = (VOID *) AlignedMemory;
+  *HostAddress = Memory;
-  *DeviceAddress = (EFI_PHYSICAL_ADDRESS) AlignedDeviceMemory;
+  *DeviceAddress = DeviceMemory;
 
   return EFI_SUCCESS;
 }

MdeModulePkg/Bus/Pci/XhciPei/XhcPeim.h

@@ -342,4 +342,32 @@ IoMmuFreeBuffer (
   IN VOID                   *Mapping
   );
 
+/**
+  Allocates aligned pages that are suitable for an OperationBusMasterCommonBuffer or
+  OperationBusMasterCommonBuffer64 mapping.
+
+  @param  Pages                 The number of pages to allocate.
+  @param  Alignment             The requested alignment of the allocation.  Must be a power of two.
+  @param  HostAddress           A pointer to store the base system memory address of the
+                                allocated range.
+  @param  DeviceAddress         The resulting map address for the bus master PCI controller to use to
+                                access the hosts HostAddress.
+  @param  Mapping               A resulting value to pass to Unmap().
+
+  @retval EFI_SUCCESS           The requested memory pages were allocated.
+  @retval EFI_UNSUPPORTED       Attributes is unsupported. The only legal attribute bits are
+                                MEMORY_WRITE_COMBINE and MEMORY_CACHED.
+  @retval EFI_INVALID_PARAMETER One or more parameters are invalid.
+  @retval EFI_OUT_OF_RESOURCES  The memory pages could not be allocated.
+
+**/
+EFI_STATUS
+IoMmuAllocateAlignedBuffer (
+  IN UINTN                  Pages,
+  IN UINTN                  Alignment,
+  OUT VOID                  **HostAddress,
+  OUT EFI_PHYSICAL_ADDRESS  *DeviceAddress,
+  OUT VOID                  **Mapping
+  );
+
 #endif