[PATCH 2/5] ArmPkg/ArmMmuLib AARCH64: get rid of needless TLB invalidation

[Ard Biesheuvel <ard.biesheuvel@linaro.org>]

Currently, we always invalidate the TLBs entirely after making
any modification to the page tables. Now that we have introduced
strict memory permissions in quite a number of places, such
modifications occur much more often, and it is better for performance
to flush only those TLB entries that are actually affected by
the changes.

Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Ard Biesheuvel <ard.biesheuvel@linaro.org>
---
 ArmPkg/Include/Library/ArmMmuLib.h                       |  3 ++-
 ArmPkg/Library/ArmLib/AArch64/ArmLibSupport.S            |  6 +++---
 ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c         | 16 +++++++---------
 ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S | 14 ++++++++------
 4 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/ArmPkg/Include/Library/ArmMmuLib.h b/ArmPkg/Include/Library/ArmMmuLib.h
index fb7fd006417c..d2725810f1c6 100644
--- a/ArmPkg/Include/Library/ArmMmuLib.h
+++ b/ArmPkg/Include/Library/ArmMmuLib.h
@@ -59,7 +59,8 @@ VOID
 EFIAPI
 ArmReplaceLiveTranslationEntry (
   IN  UINT64  *Entry,
-  IN  UINT64  Value
+  IN  UINT64  Value,
+  IN  UINT64  Address
   );
 
 EFI_STATUS
diff --git a/ArmPkg/Library/ArmLib/AArch64/ArmLibSupport.S b/ArmPkg/Library/ArmLib/AArch64/ArmLibSupport.S
index b7173e00b039..175fb58206b6 100644
--- a/ArmPkg/Library/ArmLib/AArch64/ArmLibSupport.S
+++ b/ArmPkg/Library/ArmLib/AArch64/ArmLibSupport.S
@@ -124,15 +124,15 @@ ASM_FUNC(ArmSetMAIR)
 //  IN VOID  *MVA                    // X1
 //  );
 ASM_FUNC(ArmUpdateTranslationTableEntry)
-   dc      civac, x0             // Clean and invalidate data line
-   dsb     sy
+   dsb     nshst
+   lsr     x1, x1, #12
    EL1_OR_EL2_OR_EL3(x0)
 1: tlbi    vaae1, x1             // TLB Invalidate VA , EL1
    b       4f
 2: tlbi    vae2, x1              // TLB Invalidate VA , EL2
    b       4f
 3: tlbi    vae3, x1              // TLB Invalidate VA , EL3
-4: dsb     sy
+4: dsb     nsh
    isb
    ret
 
diff --git a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
index d66df3e17a02..e1fabfcbea14 100644
--- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
+++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibCore.c
@@ -129,13 +129,14 @@ STATIC
 VOID
 ReplaceLiveEntry (
   IN  UINT64  *Entry,
-  IN  UINT64  Value
+  IN  UINT64  Value,
+  IN  UINT64  Address
   )
 {
   if (!ArmMmuEnabled ()) {
     *Entry = Value;
   } else {
-    ArmReplaceLiveTranslationEntry (Entry, Value);
+    ArmReplaceLiveTranslationEntry (Entry, Value, Address);
   }
 }
 
@@ -296,7 +297,8 @@ GetBlockEntryListFromAddress (
 
         // Fill the BlockEntry with the new TranslationTable
         ReplaceLiveEntry (BlockEntry,
-          ((UINTN)TranslationTable & TT_ADDRESS_MASK_DESCRIPTION_TABLE) | TableAttributes | TT_TYPE_TABLE_ENTRY);
+          (UINTN)TranslationTable | TableAttributes | TT_TYPE_TABLE_ENTRY,
+          RegionStart);
       }
     } else {
       if (IndexLevel != PageLevel) {
@@ -375,6 +377,8 @@ UpdateRegionMapping (
       *BlockEntry &= BlockEntryMask;
       *BlockEntry |= (RegionStart & TT_ADDRESS_MASK_BLOCK_ENTRY) | Attributes | Type;
 
+      ArmUpdateTranslationTableEntry (BlockEntry, (VOID *)RegionStart);
+
       // Go to the next BlockEntry
       RegionStart += BlockEntrySize;
       RegionLength -= BlockEntrySize;
@@ -487,9 +491,6 @@ ArmSetMemoryAttributes (
     return Status;
   }
 
-  // Invalidate all TLB entries so changes are synced
-  ArmInvalidateTlb ();
-
   return EFI_SUCCESS;
 }
 
@@ -512,9 +513,6 @@ SetMemoryRegionAttribute (
     return Status;
   }
 
-  // Invalidate all TLB entries so changes are synced
-  ArmInvalidateTlb ();
-
   return EFI_SUCCESS;
 }
 
diff --git a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S
index 90192df24f55..d40c19b2e3e5 100644
--- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S
+++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S
@@ -32,11 +32,12 @@
   dmb   sy
   dc    ivac, x0
 
-  // flush the TLBs
+  // flush translations for the target address from the TLBs
+  lsr   x2, x2, #12
   .if   \el == 1
-  tlbi  vmalle1
+  tlbi  vaae1, x2
   .else
-  tlbi  alle\el
+  tlbi  vae\el, x2
   .endif
   dsb   sy
 
@@ -48,12 +49,13 @@
 //VOID
 //ArmReplaceLiveTranslationEntry (
 //  IN  UINT64  *Entry,
-//  IN  UINT64  Value
+//  IN  UINT64  Value,
+//  IN  UINT64  Address
 //  )
 ASM_FUNC(ArmReplaceLiveTranslationEntry)
 
   // disable interrupts
-  mrs   x2, daif
+  mrs   x4, daif
   msr   daifset, #0xf
   isb
 
@@ -69,7 +71,7 @@ ASM_FUNC(ArmReplaceLiveTranslationEntry)
   b     4f
 3:__replace_entry 3
 
-4:msr   daif, x2
+4:msr   daif, x4
   ret
 
 ASM_GLOBAL ASM_PFX(ArmReplaceLiveTranslationEntrySize)
-- 
2.20.1