Re: [Patch 5/5] Vlv2TbltDevicePkg: Sync FLASH libraries from UDK2017 branch

["Guo, Mang" <mang.guo@intel.com>]

Reviewed-by:  Guo  Mang <mang.guo@intel.com>

-----Original Message-----
From: Kinney, Michael D 
Sent: Wednesday, March 28, 2018 1:34 PM
To: edk2-devel@lists.01.org
Cc: Wei, David; Guo, Mang; Yao, Jiewen; Kinney, Michael D
Subject: [Patch 5/5] Vlv2TbltDevicePkg: Sync FLASH libraries from UDK2017 branch

https://bugzilla.tianocore.org/show_bug.cgi?id=911

Update Minnow Max FLASH libraries to match libraries in the UDK2017 development branch in edk2-platforms.

https://github.com/tianocore/edk2-platforms/tree/devel-MinnowBoardMax-UDK2017

Cc: David Wei <david.wei@intel.com>
Cc: Mang Guo <mang.guo@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Michael D Kinney <michael.d.kinney@intel.com>
---
 .../PlatformFlashAccessLib.c                       | 491 ++++++++++++++++++++-
 .../PlatformFlashAccessLib.inf                     |  19 +-
 .../Library/FlashDeviceLib/FlashDeviceLib.c        | 158 ++++++-
 .../Library/FlashDeviceLib/FlashDeviceLib.inf      |   6 +-
 .../Library/FlashDeviceLib/SpiChipDefinitions.h    |   1 +
 5 files changed, 639 insertions(+), 36 deletions(-)

diff --git a/Vlv2TbltDevicePkg/Feature/Capsule/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.c b/Vlv2TbltDevicePkg/Feature/Capsule/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.c
index f3cb31daaa..9162e025ed 100644
--- a/Vlv2TbltDevicePkg/Feature/Capsule/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.c
+++ b/Vlv2TbltDevicePkg/Feature/Capsule/Library/PlatformFlashAccessLib/P
+++ latformFlashAccessLib.c
@@ -1,7 +1,7 @@
 /** @file
   Platform Flash Access library.
 
-  Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 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 @@ -11,6 +11,7 @@
   WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 **/
+#include <Uefi.h>
 
 #include <PiDxe.h>
 
@@ -19,14 +20,366 @@
 #include <Library/DebugLib.h>
 #include <Library/PcdLib.h>
 #include <Library/PlatformFlashAccessLib.h>
-#include <Library/FlashDeviceLib.h>
+//#include <Library/FlashDeviceLib.h>
 #include <Library/MemoryAllocationLib.h>
+#include <Protocol/Spi.h>
+#include <Library/CacheMaintenanceLib.h> #include "PchAccess.h"
+#include <Library/IoLib.h>
+#include <Library/UefiLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/PrintLib.h>
 
-#define SECTOR_SIZE_64KB  0x10000      // Common 64kBytes sector size
-#define ALINGED_SIZE  SECTOR_SIZE_64KB
+//#define SECTOR_SIZE_64KB  0x10000      // Common 64kBytes sector size
+//#define ALINGED_SIZE  SECTOR_SIZE_64KB
+
+#define BLOCK_SIZE 0x1000
+#define ALINGED_SIZE BLOCK_SIZE
+
+#define R_PCH_LPC_BIOS_CNTL                       0xDC
+#define B_PCH_LPC_BIOS_CNTL_SMM_BWP               0x20            ///< SMM BIOS write protect disable
+
+//
+// Prefix Opcode Index on the host SPI controller // typedef enum {
+  SPI_WREN,             // Prefix Opcode 0: Write Enable
+  SPI_EWSR,             // Prefix Opcode 1: Enable Write Status Register
+} PREFIX_OPCODE_INDEX;
+//
+// Opcode Menu Index on the host SPI controller // typedef enum {
+  SPI_READ_ID,        // Opcode 0: READ ID, Read cycle with address
+  SPI_READ,           // Opcode 1: READ, Read cycle with address
+  SPI_RDSR,           // Opcode 2: Read Status Register, No address
+  SPI_WRDI_SFDP,      // Opcode 3: Write Disable or Discovery Parameters, No address
+  SPI_SERASE,         // Opcode 4: Sector Erase (4KB), Write cycle with address
+  SPI_BERASE,         // Opcode 5: Block Erase (32KB), Write cycle with address
+  SPI_PROG,           // Opcode 6: Byte Program, Write cycle with address
+  SPI_WRSR,           // Opcode 7: Write Status Register, No address
+} SPI_OPCODE_INDEX;
 
 STATIC EFI_PHYSICAL_ADDRESS     mInternalFdAddress;
 
+EFI_SPI_PROTOCOL  *mSpiProtocol;
+
+/**
+  Read NumBytes bytes of data from the address specified by
+  PAddress into Buffer.
+
+  @param[in]      Address       The starting physical address of the read.
+  @param[in,out]  NumBytes      On input, the number of bytes to read. On output, the number
+                                of bytes actually read.
+  @param[out]     Buffer        The destination data buffer for the read.
+
+  @retval         EFI_SUCCESS       Opertion is successful.
+  @retval         EFI_DEVICE_ERROR  If there is any device errors.
+
+**/
+EFI_STATUS
+EFIAPI
+SpiFlashRead (
+  IN     UINTN     Address,
+  IN OUT UINT32    *NumBytes,
+     OUT UINT8     *Buffer
+  )
+{
+  EFI_STATUS    Status = EFI_SUCCESS;
+  UINTN         Offset = 0;
+
+  ASSERT ((NumBytes != NULL) && (Buffer != NULL));
+
+
+  //if (Address >= (UINTN)PcdGet32 (PcdGbeRomBase) && Address < (UINTN)PcdGet32 (PcdPDRRomBase)) {
+    Offset = Address - (UINTN)PcdGet32 (PcdFlashChipBase);
+
+    Status = mSpiProtocol->Execute (
+                               mSpiProtocol,
+                               1, //SPI_READ,
+                               0, //SPI_WREN,
+                               TRUE,
+                               TRUE,
+                               FALSE,
+                               Offset,
+                               BLOCK_SIZE,
+                               Buffer,
+                               EnumSpiRegionAll
+                               );
+    return Status;
+}
+
+/**
+  Write NumBytes bytes of data from Buffer to the address specified by
+  PAddresss.
+
+  @param[in]      Address         The starting physical address of the write.
+  @param[in,out]  NumBytes        On input, the number of bytes to write. On output,
+                                  the actual number of bytes written.
+  @param[in]      Buffer          The source data buffer for the write.
+
+  @retval         EFI_SUCCESS       Opertion is successful.
+  @retval         EFI_DEVICE_ERROR  If there is any device errors.
+
+**/
+EFI_STATUS
+EFIAPI
+SpiFlashWrite (
+  IN     UINTN     Address,
+  IN OUT UINT32    *NumBytes,
+  IN     UINT8     *Buffer
+  )
+{
+  EFI_STATUS                Status;
+  UINTN                     Offset;
+  UINT32                    Length;
+  UINT32                    RemainingBytes;
+
+  ASSERT ((NumBytes != NULL) && (Buffer != NULL));  ASSERT (Address >= 
+ (UINTN)PcdGet32 (PcdFlashChipBase));
+
+  Offset    = Address - (UINTN)PcdGet32 (PcdFlashChipBase);
+
+  ASSERT ((*NumBytes + Offset) <= (UINTN)PcdGet32 (PcdFlashChipSize));
+
+  Status = EFI_SUCCESS;
+  RemainingBytes = *NumBytes;
+
+  while (RemainingBytes > 0) {
+    if (RemainingBytes > SIZE_4KB) {
+      Length = SIZE_4KB;
+    } else {
+      Length = RemainingBytes;
+    }
+    Status = mSpiProtocol->Execute (
+                             mSpiProtocol,
+                             SPI_PROG,
+                             SPI_WREN,
+                             TRUE,
+                             TRUE,
+                             TRUE,
+                             (UINT32) Offset,
+                             Length,
+                             Buffer,
+                             EnumSpiRegionAll
+                             );
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+    RemainingBytes -= Length;
+    Offset += Length;
+    Buffer += Length;
+  }
+
+  //
+  // Actual number of bytes written
+  //
+  *NumBytes -= RemainingBytes;
+
+  return Status;
+}
+
+
+EFI_STATUS
+InternalReadBlock (
+  IN  EFI_PHYSICAL_ADDRESS  BaseAddress,
+  OUT VOID                  *ReadBuffer
+  )
+{
+  EFI_STATUS    Status;
+  UINT32        BlockSize;
+
+  BlockSize = BLOCK_SIZE;
+
+  Status = SpiFlashRead ((UINTN) BaseAddress, &BlockSize, ReadBuffer);
+
+  return Status;
+}
+
+/**
+  Erase the block starting at Address.
+
+  @param[in]  Address         The starting physical address of the block to be erased.
+                              This library assume that caller garantee that the PAddress
+                              is at the starting address of this block.
+  @param[in]  NumBytes        On input, the number of bytes of the logical block to be erased.
+                              On output, the actual number of bytes erased.
+
+  @retval     EFI_SUCCESS.      Opertion is successful.
+  @retval     EFI_DEVICE_ERROR  If there is any device errors.
+
+**/
+EFI_STATUS
+EFIAPI
+SpiFlashBlockErase (
+  IN UINTN    Address,
+  IN UINTN    *NumBytes
+  )
+{
+  EFI_STATUS          Status;
+  UINTN               Offset;
+  UINTN               RemainingBytes;
+
+  ASSERT (NumBytes != NULL);
+  ASSERT (Address >= (UINTN)PcdGet32 (PcdFlashChipBase));
+
+  Offset    = Address - (UINTN)PcdGet32 (PcdFlashChipBase);
+
+  ASSERT ((*NumBytes % SIZE_4KB) == 0);  ASSERT ((*NumBytes + Offset) 
+ <= (UINTN)PcdGet32 (PcdFlashChipSize));
+
+  Status = EFI_SUCCESS;
+  RemainingBytes = *NumBytes;
+
+  //
+  // To adjust the Offset with Bios/Gbe
+  //
+//  if (Address >= (UINTN)PcdGet32 (PcdFlashChipBase)) {
+//    Offset = Address - (UINTN)PcdGet32 (PcdFlashChipBase);
+
+    while (RemainingBytes > 0) {
+      Status = mSpiProtocol->Execute (
+                               mSpiProtocol,
+                               SPI_SERASE,
+                               SPI_WREN,
+                               FALSE,
+                               TRUE,
+                               FALSE,
+                               (UINT32) Offset,
+                               0,
+                               NULL,
+                               EnumSpiRegionAll
+                               );
+      if (EFI_ERROR (Status)) {
+        break;
+      }
+      RemainingBytes -= SIZE_4KB;
+      Offset         += SIZE_4KB;
+    }
+//  }
+
+  //
+  // Actual number of bytes erased
+  //
+  *NumBytes -= RemainingBytes;
+
+  return Status;
+}
+
+/**
+
+Routine Description:
+
+  Erase the whole block.
+
+Arguments:
+
+  BaseAddress  - Base address of the block to be erased.
+
+Returns:
+
+  EFI_SUCCESS - The command completed successfully.
+  Other       - Device error or wirte-locked, operation failed.
+
+**/
+EFI_STATUS
+InternalEraseBlock (
+  IN  EFI_PHYSICAL_ADDRESS BaseAddress
+  )
+{
+  EFI_STATUS                              Status;
+  UINTN                                   NumBytes;
+
+  NumBytes = BLOCK_SIZE;
+
+  Status = SpiFlashBlockErase ((UINTN) BaseAddress, &NumBytes);
+
+  return Status;
+}
+
+EFI_STATUS
+InternalCompareBlock (
+  IN  EFI_PHYSICAL_ADDRESS        BaseAddress,
+  IN  UINT8                       *Buffer
+  )
+{
+  EFI_STATUS                              Status;
+  VOID                                    *CompareBuffer;
+  UINT32                                  NumBytes;
+  INTN                                    CompareResult;
+
+  NumBytes = BLOCK_SIZE;
+  CompareBuffer = AllocatePool (NumBytes);  if (CompareBuffer == NULL) 
+ {
+    Status = EFI_OUT_OF_RESOURCES;
+    goto Done;
+  }
+
+  Status = SpiFlashRead ((UINTN) BaseAddress, &NumBytes, 
+ CompareBuffer);  if (EFI_ERROR (Status)) {
+    goto Done;
+  }
+  CompareResult = CompareMem (CompareBuffer, Buffer, BLOCK_SIZE);  if 
+ (CompareResult != 0) {
+    Status = EFI_VOLUME_CORRUPTED;
+  }
+
+Done:
+  if (CompareBuffer != NULL) {
+    FreePool (CompareBuffer);
+  }
+
+  return Status;
+}
+
+/**
+
+Routine Description:
+
+  Write a block of data.
+
+Arguments:
+
+  BaseAddress  - Base address of the block.
+  Buffer       - Data buffer.
+  BufferSize   - Size of the buffer.
+
+Returns:
+
+  EFI_SUCCESS           - The command completed successfully.
+  EFI_INVALID_PARAMETER - Invalid parameter, can not proceed.
+  Other                 - Device error or wirte-locked, operation failed.
+
+**/
+EFI_STATUS
+InternalWriteBlock (
+  IN  EFI_PHYSICAL_ADDRESS        BaseAddress,
+  IN  UINT8                       *Buffer,
+  IN  UINT32                      BufferSize
+  )
+{
+  EFI_STATUS                              Status;
+
+  Status = SpiFlashWrite ((UINTN) BaseAddress, &BufferSize, Buffer);
+
+  if (EFI_ERROR (Status)) {
+    DEBUG((DEBUG_ERROR, "\nFlash write error."));
+    return Status;
+  }
+
+  WriteBackInvalidateDataCacheRange ((VOID *) (UINTN) BaseAddress, 
+ BLOCK_SIZE);
+
+  Status = InternalCompareBlock (BaseAddress, Buffer);  if (EFI_ERROR 
+ (Status)) {
+    DEBUG((DEBUG_ERROR, "\nError when writing to BaseAddress %x with 
+ different at offset %x.", BaseAddress, Status));  } else {
+    DEBUG((DEBUG_INFO, "\nVerified data written to Block at %x is 
+ correct.", BaseAddress));  }
+
+  return Status;
+
+}
+
 /**
   Perform flash write opreation.
 
@@ -51,36 +404,122 @@ PerformFlashWrite (
   IN UINTN                        Length
   )
 {
-  EFI_STATUS          Status;
+  EFI_STATUS            Status = EFI_SUCCESS;
+  UINTN               Index;
+  EFI_PHYSICAL_ADDRESS  Address;
+  UINTN                 CountOfBlocks;
+  EFI_TPL               OldTpl;
+  BOOLEAN               FlashError;
+  UINT8                 *Buf;
+  UINTN                 LpcBaseAddress;
+  UINT8                 Data8Or;
+  UINT8                 Data8And;
+  UINT8                 BiosCntl;
 
-  DEBUG((DEBUG_INFO, "PerformFlashWrite - 0x%x(%x) - 0x%x\n", (UINTN)FlashAddress, (UINTN)FlashAddressType, Length));
+  Index             = 0;
+  Address           = 0;
+  CountOfBlocks     = 0;
+  FlashError        = FALSE;
+  Buf               = Buffer;
+
+  DEBUG((DEBUG_INFO | DEBUG_ERROR, "PerformFlashWrite - 0x%x(%x) - 
+ 0x%x\n", (UINTN)FlashAddress, (UINTN)FlashAddressType, Length));
   if (FlashAddressType == FlashAddressTypeRelativeAddress) {
     FlashAddress = FlashAddress + mInternalFdAddress;
   }
 
-  DEBUG((DEBUG_INFO, "                  - 0x%x(%x) - 0x%x\n", (UINTN)FlashAddress, (UINTN)FlashAddressType, Length));
-  LibFvbFlashDeviceBlockLock((UINTN)FlashAddress, Length, FALSE);
+  CountOfBlocks = (UINTN) (Length / BLOCK_SIZE);  Address = 
+ FlashAddress;
 
+  LpcBaseAddress = MmPciAddress (0,
+                    DEFAULT_PCI_BUS_NUMBER_PCH,
+                    PCI_DEVICE_NUMBER_PCH_LPC,
+                    PCI_FUNCTION_NUMBER_PCH_LPC,
+                    0
+                    );
+  BiosCntl = MmioRead8 (LpcBaseAddress + R_PCH_LPC_BIOS_CNTL);  if 
+ ((BiosCntl & B_PCH_LPC_BIOS_CNTL_SMM_BWP) == B_PCH_LPC_BIOS_CNTL_SMM_BWP) {
+    ///
+    /// Clear SMM_BWP bit (D31:F0:RegDCh[5])
+    ///
+    Data8And  = (UINT8) ~B_PCH_LPC_BIOS_CNTL_SMM_BWP;
+    Data8Or   = 0x00;
+
+    MmioAndThenOr8 (
+      LpcBaseAddress + R_PCH_LPC_BIOS_CNTL,
+      Data8And,
+      Data8Or
+      );
+    DEBUG((DEBUG_INFO, "PerformFlashWrite Clear SMM_BWP bit\n"));  }
+
+    //
+    // Raise TPL to TPL_NOTIFY to block any event handler,
+    // while still allowing RaiseTPL(TPL_NOTIFY) within
+    // output driver during Print()
   //
-  // Erase & Write
-  //
-  Status = LibFvbFlashDeviceBlockErase((UINTN)FlashAddress, Length);
-  ASSERT_EFI_ERROR(Status);
+    OldTpl = gBS->RaiseTPL (TPL_NOTIFY);
+    for (Index = 0; Index < CountOfBlocks; Index++) {
+      //
+      // Handle block based on address and contents.
+      //
+      if (!EFI_ERROR (InternalCompareBlock (Address, Buf))) {
+        DEBUG((DEBUG_INFO, "Skipping block at 0x%lx (already programmed)\n", Address));
+      } else {
+        //
+        // Display a dot for each block being updated.
+        //
+        Print (L".");
+
+        //
+        // Make updating process uninterruptable,
+        // so that the flash memory area is not accessed by other entities
+        // which may interfere with the updating process
+        //
+        Status  = InternalEraseBlock (Address);
   if (EFI_ERROR(Status)) {
-    LibFvbFlashDeviceBlockLock((UINTN)FlashAddress, Length, TRUE);
-    DEBUG((DEBUG_ERROR, "Flash Erase error\n"));
-    return Status;
+          gBS->RestoreTPL (OldTpl);
+          FlashError = TRUE;
+          goto Done;
   }
-
-  Status = LibFvbFlashDeviceWrite((UINTN)FlashAddress, &Length, Buffer);
-  ASSERT_EFI_ERROR(Status);
+        Status = InternalWriteBlock (
+                  Address,
+                  Buf,
+                  (UINT32)(Length > BLOCK_SIZE ? BLOCK_SIZE : Length)
+                  );
   if (EFI_ERROR(Status)) {
-    LibFvbFlashDeviceBlockLock((UINTN)FlashAddress, Length, TRUE);
-    DEBUG((DEBUG_ERROR, "Flash write error\n"));
-    return Status;
+          gBS->RestoreTPL (OldTpl);
+          FlashError = TRUE;
+          goto Done;
+        }
   }
 
-  LibFvbFlashDeviceBlockLock((UINTN)FlashAddress, Length, TRUE);
+      //
+      // Move to next block to update.
+      //
+      Address += BLOCK_SIZE;
+      Buf += BLOCK_SIZE;
+      if (Length > BLOCK_SIZE) {
+        Length -= BLOCK_SIZE;
+      } else {
+        Length = 0;
+      }
+    }
+    gBS->RestoreTPL (OldTpl);
+
+  Done:
+  if ((BiosCntl & B_PCH_LPC_BIOS_CNTL_SMM_BWP) == B_PCH_LPC_BIOS_CNTL_SMM_BWP) {
+    //
+    // Restore original control setting
+    //
+    MmioWrite8 (LpcBaseAddress + R_PCH_LPC_BIOS_CNTL, BiosCntl);
+    }
+
+  //
+  // Print flash update failure message if error detected.
+  //
+  if (FlashError) {
+    Print (L"No %r\n", Status);
+  }
 
   return EFI_SUCCESS;
 }
@@ -183,8 +622,16 @@ PerformFlashAccessLibConstructor (
   VOID
   )
 {
+  EFI_STATUS Status;
   mInternalFdAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)PcdGet32(PcdFlashAreaBaseAddress);
   DEBUG((DEBUG_INFO, "PcdFlashAreaBaseAddress - 0x%x\n", mInternalFdAddress));
 
+  Status = gBS->LocateProtocol (
+                  &gEfiSpiProtocolGuid,
+                  NULL,
+                  (VOID **) &mSpiProtocol
+                  );
+  ASSERT_EFI_ERROR(Status);
+
   return EFI_SUCCESS;
 }
diff --git a/Vlv2TbltDevicePkg/Feature/Capsule/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.inf b/Vlv2TbltDevicePkg/Feature/Capsule/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.inf
index 17ab38e27f..fbbdb91b64 100644
--- a/Vlv2TbltDevicePkg/Feature/Capsule/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.inf
+++ b/Vlv2TbltDevicePkg/Feature/Capsule/Library/PlatformFlashAccessLib/P
+++ latformFlashAccessLib.inf
@@ -1,7 +1,7 @@
 ## @file
 #  Platform Flash Access library.
 #
-#  Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>
+#  Copyright (c) 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 @@ -36,13 +36,26 @@ [Packages]
   MdeModulePkg/MdeModulePkg.dec
   SignedCapsulePkg/SignedCapsulePkg.dec
   Vlv2TbltDevicePkg/PlatformPkg.dec
+  Vlv2DeviceRefCodePkg/Vlv2DeviceRefCodePkg.dec
 
 [LibraryClasses]
   BaseMemoryLib
+  IoLib
   PcdLib
   DebugLib
-  FlashDeviceLib
+#  FlashDeviceLib
   MemoryAllocationLib
+  CacheMaintenanceLib
+
+[Guids]
+  gEdkiiSystemFmpCapsuleConfigFileGuid          ## SOMETIMES_CONSUMES ## GUID
+
+[Protocols]
+  gEfiSpiProtocolGuid                          ## CONSUMES
 
 [Pcd]
-  gPlatformModuleTokenSpaceGuid.PcdFlashAreaBaseAddress
+  gPlatformModuleTokenSpaceGuid.PcdFlashAreaBaseAddress  ## SOMETIMES_CONSUMES
+  gPlatformModuleTokenSpaceGuid.PcdFlashChipBase         ## SOMETIMES_CONSUMES
+  gPlatformModuleTokenSpaceGuid.PcdFlashChipSize         ## SOMETIMES_CONSUMES
+  gEfiMdePkgTokenSpaceGuid.PcdPciExpressBaseAddress      ## SOMETIMES_CONSUMES
+
diff --git a/Vlv2TbltDevicePkg/Library/FlashDeviceLib/FlashDeviceLib.c b/Vlv2TbltDevicePkg/Library/FlashDeviceLib/FlashDeviceLib.c
index 8e6e5b4728..392ad1c14b 100644
--- a/Vlv2TbltDevicePkg/Library/FlashDeviceLib/FlashDeviceLib.c
+++ b/Vlv2TbltDevicePkg/Library/FlashDeviceLib/FlashDeviceLib.c
@@ -19,13 +19,14 @@
 #include <Library/FlashDeviceLib.h>
 #include <Library/DebugLib.h>
 #include <Library/BaseLib.h>
+#include <Library/UefiBootServicesTableLib.h>
+#include <Library/UefiRuntimeServicesTableLib.h>
 #include <Library/BaseMemoryLib.h>
+#include <Library/UefiRuntimeLib.h>
+#include <Protocol/SmmBase2.h>
 #include <Guid/EventGroup.h>
-#include <Library/SpiFlash.H>
+#include "SpiChipDefinitions.h"
 
-#define FLASH_SIZE  0x400000
-
-#define FLASH_DEVICE_BASE_ADDRESS (0xFFFFFFFF-FLASH_SIZE+1)  UINTN FlashDeviceBase = FLASH_DEVICE_BASE_ADDRESS;
 
 EFI_SPI_PROTOCOL *mSpiProtocol = NULL;
@@ -77,11 +78,11 @@ SpiFlashBlockErase (
   UINT32              SpiAddress;
 
   SpiAddress = (UINT32)(UINTN)(BaseAddress) - (UINT32)FlashDeviceBase;
-  SectorSize = SECTOR_SIZE_64KB;
+  SectorSize = SECTOR_SIZE_4KB;
   while ( (NumBytes > 0) && (NumBytes <= MAX_FWH_SIZE) ) {
     Status = mSpiProtocol->Execute (
                              mSpiProtocol,
-                             SPI_BERASE,
+                             SPI_SERASE,
                              SPI_WREN,
                              FALSE,
                              TRUE,
@@ -319,3 +320,148 @@ LibFvbFlashDeviceBlockLock (
   return Status;
 }
 
+VOID
+EFIAPI
+LibFvbFlashDeviceVirtualAddressChangeNotifyEvent (
+  IN EFI_EVENT        Event,
+  IN VOID             *Context
+  )
+{
+  gRT->ConvertPointer (0, (VOID **) &mSpiProtocol);
+  gRT->ConvertPointer (0, (VOID **) &FlashDeviceBase); }
+
+
+/**
+  The library constructuor.
+
+  The function does the necessary initialization work for this library  
+ instance. Please put all initialization works in it.
+
+  @param[in]  ImageHandle       The firmware allocated handle for the UEFI image.
+  @param[in]  SystemTable       A pointer to the EFI system table.
+
+  @retval     EFI_SUCCESS       The function always return EFI_SUCCESS for now.
+                                It will ASSERT on error for debug version.
+  @retval     EFI_ERROR         Please reference LocateProtocol for error code details.
+
+**/
+EFI_STATUS
+EFIAPI
+LibFvbFlashDeviceSupportInit (
+  IN EFI_HANDLE         ImageHandle,
+  IN EFI_SYSTEM_TABLE   *SystemTable
+  )
+{
+  EFI_STATUS Status;
+  EFI_EVENT  Event;
+  UINT8                         SfId[3];
+  UINT8                         FlashIndex;
+  UINT8                         SpiReadError;
+  UINT8                         SpiNotMatchError;
+  EFI_SMM_BASE2_PROTOCOL       *SmmBase;
+  BOOLEAN                       InSmm;
+
+  SpiReadError     = 0x00;
+  SpiNotMatchError = 0x00;
+
+  InSmm = FALSE;
+  Status = gBS->LocateProtocol (
+                  &gEfiSmmBase2ProtocolGuid,
+                  NULL,
+                  (void **)&SmmBase
+                  );
+  if (!EFI_ERROR(Status)) {
+    Status = SmmBase->InSmm(SmmBase, &InSmm);
+    if (EFI_ERROR(Status)) {
+      InSmm = FALSE;
+    }
+  }
+
+  if (!InSmm) {
+    Status = gBS->LocateProtocol (
+                  &gEfiSpiProtocolGuid,
+                  NULL,
+                  (VOID **)&mSpiProtocol
+                  );
+    ASSERT_EFI_ERROR (Status);
+
+    Status = gBS->CreateEventEx (
+                  EVT_NOTIFY_SIGNAL,
+                  TPL_NOTIFY,
+                  LibFvbFlashDeviceVirtualAddressChangeNotifyEvent,
+                  NULL,
+                  &gEfiEventVirtualAddressChangeGuid,
+                  &Event
+                  );
+    ASSERT_EFI_ERROR (Status);
+  } else {
+    Status = gBS->LocateProtocol (
+                    &gEfiSmmSpiProtocolGuid,
+                    NULL,
+                    (VOID **)&mSpiProtocol
+                    );
+    ASSERT_EFI_ERROR (Status);
+  }
+
+
+  for (FlashIndex = EnumSpiFlashW25Q64; FlashIndex < EnumSpiFlashMax; FlashIndex++) {
+    Status = mSpiProtocol->Init (mSpiProtocol, &(mInitTable[FlashIndex]));
+    if (!EFI_ERROR (Status)) {
+      //
+      // Read Vendor/Device IDs to check if the driver supports the Serial Flash device.
+      //
+      Status = mSpiProtocol->Execute (
+                               mSpiProtocol,
+                               SPI_READ_ID,
+                               SPI_WREN,
+                               TRUE,
+                               FALSE,
+                               FALSE,
+                               0,
+                               3,
+                               SfId,
+                               EnumSpiRegionAll
+                               );
+      if (!EFI_ERROR (Status)) {
+        if ((SfId[0] == mInitTable[FlashIndex].VendorId)  &&
+            (SfId[1] == mInitTable[FlashIndex].DeviceId0) &&
+            (SfId[2] == mInitTable[FlashIndex].DeviceId1)) {
+            //
+            // Found a matching SPI device, FlashIndex now contains flash device.
+            //
+            DEBUG ((DEBUG_ERROR, "OK - Found SPI Flash Type in SPI Flash Driver, Device Type ID 0 = 0x%02x!\n", mInitTable[FlashIndex].DeviceId0));
+            DEBUG ((DEBUG_ERROR, "Device Type ID 1 = 0x%02x!\n", 
+ mInitTable[FlashIndex].DeviceId1));
+
+            if (mInitTable[FlashIndex].BiosStartOffset == (UINTN) (-1)) {
+              DEBUG ((DEBUG_ERROR, "ERROR - The size of BIOS image is bigger than SPI Flash device!\n"));
+              CpuDeadLoop ();
+            }
+            break;
+        } else {
+          SpiNotMatchError++;
+        }
+      } else {
+        SpiReadError++;
+      }
+    }
+  }
+
+  DEBUG ((DEBUG_ERROR, "SPI flash chip VID = 0x%X, DID0 = 0x%X, DID1 = 
+ 0x%X\n", SfId[0], SfId[1], SfId[2]));
+
+  if (FlashIndex < EnumSpiFlashMax)  {
+    return EFI_SUCCESS;
+  } else {
+  if (SpiReadError != 0) {
+      DEBUG ((DEBUG_ERROR, "ERROR - SPI Read ID execution failed! Error Count = %d\n", SpiReadError));
+   }
+    else {
+      if (SpiNotMatchError != 0) {
+        DEBUG ((DEBUG_ERROR, "ERROR - No supported SPI flash chip found! Error Count = %d\n", SpiNotMatchError));
+        DEBUG ((DEBUG_ERROR, "SPI flash chip VID = 0x%X, DID0 = 0x%X, DID1 = 0x%X\n", SfId[0], SfId[1], SfId[2]));
+      }
+    }
+    return EFI_UNSUPPORTED;
+  }
+}
+
diff --git a/Vlv2TbltDevicePkg/Library/FlashDeviceLib/FlashDeviceLib.inf b/Vlv2TbltDevicePkg/Library/FlashDeviceLib/FlashDeviceLib.inf
index 955b2479f6..3331432b68 100644
--- a/Vlv2TbltDevicePkg/Library/FlashDeviceLib/FlashDeviceLib.inf
+++ b/Vlv2TbltDevicePkg/Library/FlashDeviceLib/FlashDeviceLib.inf
@@ -16,7 +16,7 @@
 
 [Defines]
   INF_VERSION                    = 0x00010005
-  BASE_NAME                      = FlashDeviceLibRuntimeSmm
+  BASE_NAME                      = FlashDeviceLib
   FILE_GUID                      = E38A1C3C-928C-4bf7-B6C1-7F0EF163FAA5
   MODULE_TYPE                    = DXE_DRIVER
   VERSION_STRING                 = 1.0
@@ -32,7 +32,6 @@ [Defines]
 
 [Sources]
   FlashDeviceLib.c
-  FlashDeviceLibDxeRuntimeSmm.c
 
 
 [Packages]
@@ -44,9 +43,6 @@ [Packages]
 [LibraryClasses]
   DebugLib
 
-[Guids]
-  gEfiEventVirtualAddressChangeGuid
-
 [Protocols]
   gEfiSpiProtocolGuid
   gEfiSmmSpiProtocolGuid
diff --git a/Vlv2TbltDevicePkg/Library/FlashDeviceLib/SpiChipDefinitions.h b/Vlv2TbltDevicePkg/Library/FlashDeviceLib/SpiChipDefinitions.h
index fdf742c1ba..954aadc174 100644
--- a/Vlv2TbltDevicePkg/Library/FlashDeviceLib/SpiChipDefinitions.h
+++ b/Vlv2TbltDevicePkg/Library/FlashDeviceLib/SpiChipDefinitions.h
@@ -17,6 +17,7 @@
 #include <Library/SpiFlash.H>
 
 #define FLASH_SIZE  0x400000
+#define FLASH_DEVICE_BASE_ADDRESS (0xFFFFFFFF-FLASH_SIZE+1)
 
 //
 // Serial Flash device initialization data table provided to the
--
2.14.2.windows.3