[RFC] [PATCH 2/3] MdeModulePkg/PciHostBridge: Add IOMMU support.

[Jiewen Yao <jiewen.yao@intel.com>]

The responsibility of PciHostBridge is to allocate IOMMU
page aligned memory for Map and AllocateBuffer,
because PciHostBridge driver already handles Map() request
to allocate another buffer for DMA read/write.

PciHostBridge does not set IOMMU attribute because it does
not know which device request the DMA. This work is done
by PciBus driver.

Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Cc: Leo Duran <leo.duran@amd.com>
Cc: Brijesh Singh <brijesh.singh@amd.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiewen Yao <jiewen.yao@intel.com>
---
 MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridge.c      |   3 +
 MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridgeDxe.inf |   1 +
 MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridge.h      |   7 +
 MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridgeIo.c    | 172 +++++++++++++++++++-
 4 files changed, 178 insertions(+), 5 deletions(-)

diff --git a/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridge.c b/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridge.c
index 9005dee..35233a7 100644
--- a/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridge.c
+++ b/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridge.c
@@ -28,6 +28,9 @@ GLOBAL_REMOVE_IF_UNREFERENCED CHAR16 *mPciResourceTypeStr[] = {
   L"I/O", L"Mem", L"PMem", L"Mem64", L"PMem64", L"Bus"
 };
 
+EDKII_IOMMU_PROTOCOL        *gIoMmuProtocol;
+UINTN                       mIoMmuPageSize = 1;
+
 /**
   Ensure the compatibility of an IO space descriptor with the IO aperture.
 
diff --git a/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridgeDxe.inf b/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridgeDxe.inf
index d8b0439..2d3c8c9 100644
--- a/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridgeDxe.inf
+++ b/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciHostBridgeDxe.inf
@@ -49,6 +49,7 @@
   gEfiDevicePathProtocolGuid                      ## BY_START
   gEfiPciRootBridgeIoProtocolGuid                 ## BY_START
   gEfiPciHostBridgeResourceAllocationProtocolGuid ## BY_START
+  gEdkiiIoMmuProtocolGuid                         ## CONSUMES
 
 [Depex]
   gEfiCpuIo2ProtocolGuid AND
diff --git a/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridge.h b/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridge.h
index 13185b4..4d21d10 100644
--- a/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridge.h
+++ b/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridge.h
@@ -27,6 +27,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Protocol/CpuIo2.h>
 #include <Protocol/DevicePath.h>
 #include <Protocol/PciRootBridgeIo.h>
+#include <Protocol/IoMmu.h>
 #include <Library/DebugLib.h>
 #include <Library/DevicePathLib.h>
 #include <Library/BaseMemoryLib.h>
@@ -50,6 +51,8 @@ typedef struct {
   EFI_PCI_ROOT_BRIDGE_IO_PROTOCOL_OPERATION Operation;
   UINTN                                     NumberOfBytes;
   UINTN                                     NumberOfPages;
+  UINTN                                     MappedNumberOfBytes;
+  UINTN                                     MappedNumberOfPages;
   EFI_PHYSICAL_ADDRESS                      HostAddress;
   EFI_PHYSICAL_ADDRESS                      MappedHostAddress;
 } MAP_INFO;
@@ -575,4 +578,8 @@ RootBridgeIoConfiguration (
 
 extern EFI_METRONOME_ARCH_PROTOCOL *mMetronome;
 extern EFI_CPU_IO2_PROTOCOL         *mCpuIo;
+
+extern EDKII_IOMMU_PROTOCOL        *gIoMmuProtocol;
+extern UINTN                       mIoMmuPageSize;
+
 #endif
diff --git a/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridgeIo.c b/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridgeIo.c
index 8af131b..47ea697 100644
--- a/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridgeIo.c
+++ b/MdeModulePkg/Bus/Pci/PciHostBridgeDxe/PciRootBridgeIo.c
@@ -1022,6 +1022,121 @@ RootBridgeIoPciWrite (
 }
 
 /**
+  Allocates one or more 4KB pages of a certain memory type at a specified alignment.
+
+  Allocates the number of 4KB pages specified by Pages of a certain memory type with an alignment
+  specified by Alignment.  The allocated buffer is returned.  If Pages is 0, then NULL is returned.
+  If there is not enough memory at the specified alignment remaining to satisfy the request, then
+  NULL is returned.
+  If Alignment is not a power of two and Alignment is not zero, then ASSERT().
+  If Pages plus EFI_SIZE_TO_PAGES (Alignment) overflows, then ASSERT().
+
+  @param  Type                  The type of allocation to perform.
+  @param  MemoryType            The type of memory to allocate.
+  @param  Pages                 The number of 4 KB pages to allocate.
+  @param  Alignment             The requested alignment of the allocation.  Must be a power of two.
+                                If Alignment is zero, then byte alignment is used.
+  @param  Address               Pointer to a physical address.
+
+  @return Memory Allocation Status.
+
+**/
+EFI_STATUS
+InternalAllocateAlignedPagesWithAllocateType (
+  IN EFI_ALLOCATE_TYPE         Type,
+  IN EFI_MEMORY_TYPE           MemoryType,
+  IN UINTN                     Pages,
+  IN UINTN                     Alignment,
+  IN OUT EFI_PHYSICAL_ADDRESS  *Address
+  )
+{
+  EFI_STATUS            Status;
+  EFI_PHYSICAL_ADDRESS  Memory;
+  UINTN                 AlignedMemory;
+  UINTN                 AlignmentMask;
+  UINTN                 UnalignedPages;
+  UINTN                 RealPages;
+
+  //
+  // Alignment must be a power of two or zero.
+  //
+  ASSERT ((Alignment & (Alignment - 1)) == 0);
+
+  if (Pages == 0) {
+    return EFI_INVALID_PARAMETER;
+  }
+  if (Alignment > EFI_PAGE_SIZE) {
+    //
+    // 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);
+
+    Memory         = *Address;
+    Status         = gBS->AllocatePages (Type, MemoryType, RealPages, &Memory);
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
+    AlignedMemory  = ((UINTN) Memory + AlignmentMask) & ~AlignmentMask;
+    UnalignedPages = EFI_SIZE_TO_PAGES (AlignedMemory - (UINTN) Memory);
+    if (UnalignedPages > 0) {
+      //
+      // Free first unaligned page(s).
+      //
+      Status = gBS->FreePages (Memory, UnalignedPages);
+      ASSERT_EFI_ERROR (Status);
+    }
+    Memory         = (EFI_PHYSICAL_ADDRESS) (AlignedMemory + EFI_PAGES_TO_SIZE (Pages));
+    UnalignedPages = RealPages - Pages - UnalignedPages;
+    if (UnalignedPages > 0) {
+      //
+      // Free last unaligned page(s).
+      //
+      Status = gBS->FreePages (Memory, UnalignedPages);
+      ASSERT_EFI_ERROR (Status);
+    }
+  } else {
+    //
+    // Do not over-allocate pages in this case.
+    //
+    Memory = *Address;
+    Status = gBS->AllocatePages (Type, MemoryType, Pages, &Memory);
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
+    AlignedMemory  = (UINTN) Memory;
+  }
+  *Address = AlignedMemory;
+  return EFI_SUCCESS;
+}
+
+/**
+  Return if a value is aligned.
+
+  @param Value the value to be checked
+  @param Alignment the alignment to be checked with.
+
+  @retval TRUE  The value is aligned
+  @retval FALSE The value is not aligned.
+**/
+BOOLEAN
+InternalIsAlgined (
+  IN UINTN  Value,
+  IN UINTN  Alignment
+  )
+{
+  if (Value != ALIGN_VALUE(Value, Alignment)) {
+    return FALSE;
+  } else {
+    return FALSE;
+  }
+}
+
+/**
   Provides the PCI controller-specific address needed to access
   system memory for DMA.
 
@@ -1057,6 +1172,8 @@ RootBridgeIoMap (
   PCI_ROOT_BRIDGE_INSTANCE                          *RootBridge;
   EFI_PHYSICAL_ADDRESS                              PhysicalAddress;
   MAP_INFO                                          *MapInfo;
+  BOOLEAN                                           NeedMap;
+  EFI_PHYSICAL_ADDRESS                              MaxAddress;
 
   if (HostAddress == NULL || NumberOfBytes == NULL || DeviceAddress == NULL ||
       Mapping == NULL) {
@@ -1072,12 +1189,40 @@ RootBridgeIoMap (
 
   RootBridge = ROOT_BRIDGE_FROM_THIS (This);
 
+  if (gIoMmuProtocol == NULL) {
+    gBS->LocateProtocol (
+          &gEdkiiIoMmuProtocolGuid,
+          NULL,
+          (VOID **) &gIoMmuProtocol
+          );
+    if (gIoMmuProtocol != NULL) {
+      gIoMmuProtocol->GetPageSize (gIoMmuProtocol, &mIoMmuPageSize);
+      ASSERT ((mIoMmuPageSize & (mIoMmuPageSize - 1)) == 0);
+    }
+  }
+
   PhysicalAddress = (EFI_PHYSICAL_ADDRESS) (UINTN) HostAddress;
+
+  NeedMap = FALSE;
+  MaxAddress = (UINT64)-1;
+
   if ((!RootBridge->DmaAbove4G ||
        (Operation != EfiPciOperationBusMasterRead64 &&
         Operation != EfiPciOperationBusMasterWrite64 &&
         Operation != EfiPciOperationBusMasterCommonBuffer64)) &&
       ((PhysicalAddress + *NumberOfBytes) > SIZE_4GB)) {
+    NeedMap = TRUE;
+    MaxAddress = SIZE_4GB - 1;
+  }
+
+  if (gIoMmuProtocol != NULL) {
+    if ((!InternalIsAlgined (*NumberOfBytes, mIoMmuPageSize)) ||
+        (!InternalIsAlgined ((UINTN)HostAddress, mIoMmuPageSize))) {
+      NeedMap = TRUE;
+    }
+  }
+
+  if (NeedMap) {
 
     //
     // If the root bridge or the device cannot handle performing DMA above
@@ -1113,15 +1258,18 @@ RootBridgeIoMap (
     MapInfo->NumberOfBytes     = *NumberOfBytes;
     MapInfo->NumberOfPages     = EFI_SIZE_TO_PAGES (MapInfo->NumberOfBytes);
     MapInfo->HostAddress       = PhysicalAddress;
-    MapInfo->MappedHostAddress = SIZE_4GB - 1;
+    MapInfo->MappedHostAddress = MaxAddress;
+    MapInfo->MappedNumberOfBytes = ALIGN_VALUE (MapInfo->NumberOfBytes, mIoMmuPageSize);
+    MapInfo->MappedNumberOfPages = EFI_SIZE_TO_PAGES (MapInfo->MappedNumberOfBytes);
 
     //
     // Allocate a buffer below 4GB to map the transfer to.
     //
-    Status = gBS->AllocatePages (
+    Status = InternalAllocateAlignedPagesWithAllocateType (
                     AllocateMaxAddress,
                     EfiBootServicesData,
-                    MapInfo->NumberOfPages,
+                    MapInfo->MappedNumberOfPages,
+                    mIoMmuPageSize,
                     &MapInfo->MappedHostAddress
                     );
     if (EFI_ERROR (Status)) {
@@ -1240,7 +1388,7 @@ RootBridgeIoUnmap (
   //
   // Free the mapped buffer and the MAP_INFO structure.
   //
-  gBS->FreePages (MapInfo->MappedHostAddress, MapInfo->NumberOfPages);
+  gBS->FreePages (MapInfo->MappedHostAddress, MapInfo->MappedNumberOfPages);
   FreePool (Mapping);
   return EFI_SUCCESS;
 }
@@ -1286,6 +1434,7 @@ RootBridgeIoAllocateBuffer (
   EFI_PHYSICAL_ADDRESS      PhysicalAddress;
   PCI_ROOT_BRIDGE_INSTANCE  *RootBridge;
   EFI_ALLOCATE_TYPE         AllocateType;
+  UINTN                     Size;
 
   //
   // Validate Attributes
@@ -1321,10 +1470,16 @@ RootBridgeIoAllocateBuffer (
     AllocateType    = AllocateMaxAddress;
     PhysicalAddress = (EFI_PHYSICAL_ADDRESS) (SIZE_4GB - 1);
   }
-  Status = gBS->AllocatePages (
+  if (gIoMmuProtocol != NULL) {
+    Size = EFI_PAGES_TO_SIZE(Pages);
+    Size = ALIGN_VALUE(EFI_PAGES_TO_SIZE(Pages), mIoMmuPageSize);
+    Pages = EFI_SIZE_TO_PAGES (Size);
+  }
+  Status = InternalAllocateAlignedPagesWithAllocateType (
                   AllocateType,
                   MemoryType,
                   Pages,
+                  mIoMmuPageSize,
                   &PhysicalAddress
                   );
   if (!EFI_ERROR (Status)) {
@@ -1356,6 +1511,13 @@ RootBridgeIoFreeBuffer (
   OUT VOID                             *HostAddress
   )
 {
+  UINTN                     Size;
+
+  if (gIoMmuProtocol != NULL) {
+    Size = EFI_PAGES_TO_SIZE(Pages);
+    Size = ALIGN_VALUE(EFI_PAGES_TO_SIZE(Pages), mIoMmuPageSize);
+    Pages = EFI_SIZE_TO_PAGES (Size);
+  }
   return gBS->FreePages ((EFI_PHYSICAL_ADDRESS) (UINTN) HostAddress, Pages);
 }
 
-- 
2.7.4.windows.1