[PATCH v2 2/2] Fix bug on SRIOV ReservedBusNum when ARI enable.

[Foster Nong <foster.nong@intel.com>]

If a device which support both features SR-IOV/ARI  has multi
functions, which maybe support 8-255. After enable ARI forwarding in
the root port and ARI Capable Hierarchy in the SR-IOV PF0.
The device will support and expose multi functions(0-255) with ARI ID routing.
In next device loop in below for() code, actually it still be in the
same SR-IOV device, and just some PF which is over 8 or higher
one(n*8), PciAllocateBusNumber() will allocate bus
number(ReservedBusNum - TempReservedBusNum)) for this PF. if reset
TempReservedBusNum as 0 in this case,it will allocate wrong bus number
for this PF because TempReservedBusNum should be total previous PF's
reserved bus numbers.

code:
  for (Device = 0; Device <= PCI_MAX_DEVICE; Device++) {
    TempReservedBusNum = 0;
    for (Func = 0; Func <= PCI_MAX_FUNC; Func++) {
    //
    // Check to see whether a pci device is present
    //
    Status = PciDevicePresent (
                 PciRootBridgeIo,
                 &Pci,
                 StartBusNumber,
                 Device,
                 Func
                 );
    ...
    Status = PciAllocateBusNumber (PciDevice, *SubBusNumber,
    (UINT8)(PciDevice->ReservedBusNum - TempReservedBusNum), SubBusNumber);

The solution is add a new flag IsAriEnabled to help handle this case.
if ARI is enabled, then TempReservedBusNum will not be reset again
during all functions(1-255) scan with checking flag IsAriEnabled.

Signed-off-by: Foster Nong <foster.nong@intel.com>
---
 MdeModulePkg/Bus/Pci/PciBusDxe/PciEnumeratorSupport.c |  1 +
 MdeModulePkg/Bus/Pci/PciBusDxe/PciLib.c               | 19 ++++++++++++++++++-
 MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h               |  1 +
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/MdeModulePkg/Bus/Pci/PciBusDxe/PciEnumeratorSupport.c b/MdeModulePkg/Bus/Pci/PciBusDxe/PciEnumeratorSupport.c
index bc20da1f38..8eca859695 100644
--- a/MdeModulePkg/Bus/Pci/PciBusDxe/PciEnumeratorSupport.c
+++ b/MdeModulePkg/Bus/Pci/PciBusDxe/PciEnumeratorSupport.c
@@ -2286,6 +2286,7 @@ CreatePciIoDevice (
                          &Data32
                          );
       if ((Data32 & EFI_PCIE_CAPABILITY_DEVICE_CAPABILITIES_2_ARI_FORWARDING) != 0) {
+        PciIoDevice->IsAriEnabled = TRUE;
         //
         // ARI forward support in bridge, so enable it.
         //
diff --git a/MdeModulePkg/Bus/Pci/PciBusDxe/PciLib.c b/MdeModulePkg/Bus/Pci/PciBusDxe/PciLib.c
index d5e3ef4d3f..3a57c05755 100644
--- a/MdeModulePkg/Bus/Pci/PciBusDxe/PciLib.c
+++ b/MdeModulePkg/Bus/Pci/PciBusDxe/PciLib.c
@@ -1106,6 +1106,7 @@ PciScanBus (
   EFI_PCI_ROOT_BRIDGE_IO_PROTOCOL    *PciRootBridgeIo;
   BOOLEAN                            BusPadding;
   UINT32                             TempReservedBusNum;
+  BOOLEAN                            IsAriEnabled;
 
   PciRootBridgeIo = Bridge->PciRootBridgeIo;
   SecondBus       = 0;
@@ -1116,9 +1117,12 @@ PciScanBus (
   BusPadding      = FALSE;
   PciDevice       = NULL;
   PciAddress      = 0;
+  IsAriEnabled    = FALSE;
 
   for (Device = 0; Device <= PCI_MAX_DEVICE; Device++) {
-    TempReservedBusNum = 0;
+    if (!IsAriEnabled) {
+      TempReservedBusNum = 0;
+    }
     for (Func = 0; Func <= PCI_MAX_FUNC; Func++) {
       //
       // Check to see whether a pci device is present
@@ -1157,6 +1161,19 @@ PciScanBus (
       if (EFI_ERROR (Status)) {
         continue;
       }
+      //
+      // Per Pcie spec ARI Extended Capability
+      // This capability must be implemented by each function in an ARI device.
+      // It is not applicable to a Root Port, a Switch Downstream Port, an RCiEP, or a Root Complex Event Collector
+      //
+      if (((Device == 0) && (Func == 0)) && (PciDevice->IsAriEnabled)) {
+        IsAriEnabled = TRUE;
+      }
+      if (PciDevice->IsAriEnabled != IsAriEnabled) {
+        DEBUG ((DEBUG_ERROR, "ERROR: %02x:%02x:%02x device ARI Feature(%x) is not consistent with others Function\n",
+                StartBusNumber, Device, Func, PciDevice->IsAriEnabled));
+        return EFI_DEVICE_ERROR;
+      }
 
       PciAddress = EFI_PCI_ADDRESS (StartBusNumber, Device, Func, 0);
 
diff --git a/MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h b/MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h
index 4b58c3ea9b..ca5c06204d 100644
--- a/MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h
+++ b/MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h
@@ -262,6 +262,7 @@ struct _PCI_IO_DEVICE {
   EFI_ACPI_ADDRESS_SPACE_DESCRIPTOR            *BusNumberRanges;
 
   BOOLEAN                                      IsPciExp;
+  BOOLEAN                                      IsAriEnabled;
   //
   // For SR-IOV
   //
-- 
2.37.1.windows.1