diff mbox series

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

Message ID 20190107071504.2431-3-ard.biesheuvel@linaro.org
State Accepted
Commit d5788777bcc75936cc0e6acb540a5ee6ac77866b
Headers show
Series memory/MMU hardening for AArch64 | expand

Commit Message

Ard Biesheuvel Jan. 7, 2019, 7:15 a.m. UTC
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(-)

-- 
2.20.1

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel

Comments

Leif Lindholm Jan. 23, 2019, 3:46 p.m. UTC | #1
On Mon, Jan 07, 2019 at 08:15:01AM +0100, Ard Biesheuvel wrote:
> 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);


OK, this whole patch took a few times around the loop before I think I
caught on what was happening.

I think I'm down to the only things confusing me being:
- The name "Address" to refer to something that is always the start
  address of a 4KB-aligned translation region.
  Is this because the function will be usable in other contexts in
  later patches?
- Why drop the & TT_ADDRESS_MASK_DESCRIPTION_TABLE bit here?
  Was it just always pointless and you decided to drop it while you
  were at it?

/
    Leif

>        }

>      } 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

> 

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
Ard Biesheuvel Jan. 23, 2019, 3:55 p.m. UTC | #2
On Wed, 23 Jan 2019 at 16:46, Leif Lindholm <leif.lindholm@linaro.org> wrote:
>

> On Mon, Jan 07, 2019 at 08:15:01AM +0100, Ard Biesheuvel wrote:

> > 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);

>


/me pages in the data ...

> OK, this whole patch took a few times around the loop before I think I

> caught on what was happening.

>

> I think I'm down to the only things confusing me being:

> - The name "Address" to refer to something that is always the start

>   address of a 4KB-aligned translation region.

>   Is this because the function will be usable in other contexts in

>   later patches?


I could change it to VirtualAddress if you prefer. It is only passed
for TLB maintenance which is only needed at page granularity, and the
low bits are shifted out anyway.

> - Why drop the & TT_ADDRESS_MASK_DESCRIPTION_TABLE bit here?

>   Was it just always pointless and you decided to drop it while you

>   were at it?

>


IIRC yes. It is a newly allocated page, so the masking does not do anything.


> /

>     Leif

>

> >        }

> >      } 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

> >

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
Leif Lindholm Jan. 23, 2019, 4:12 p.m. UTC | #3
On Wed, Jan 23, 2019 at 04:55:28PM +0100, Ard Biesheuvel wrote:
> On Wed, 23 Jan 2019 at 16:46, Leif Lindholm <leif.lindholm@linaro.org> wrote:

> >

> > On Mon, Jan 07, 2019 at 08:15:01AM +0100, Ard Biesheuvel wrote:

> > > 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/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);

> >

> 

> /me pages in the data ...

> 

> > OK, this whole patch took a few times around the loop before I think I

> > caught on what was happening.

> >

> > I think I'm down to the only things confusing me being:

> > - The name "Address" to refer to something that is always the start

> >   address of a 4KB-aligned translation region.

> >   Is this because the function will be usable in other contexts in

> >   later patches?

> 

> I could change it to VirtualAddress if you prefer.

> It is only passed

> for TLB maintenance which is only needed at page granularity, and the

> low bits are shifted out anyway.


Yeah, exactly. It would just be nice if the name reflected that. Not
sure VirtualAddress does. VirtualBase? PageBase?

> > - Why drop the & TT_ADDRESS_MASK_DESCRIPTION_TABLE bit here?

> >   Was it just always pointless and you decided to drop it while you

> >   were at it?

> 

> IIRC yes. It is a newly allocated page, so the masking does not do anything.


Yeah, that's fair enough.
Just made me wonder if anything unobvious was going on :)

/
    Leif
_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
Ard Biesheuvel Jan. 23, 2019, 4:16 p.m. UTC | #4
On Wed, 23 Jan 2019 at 17:13, Leif Lindholm <leif.lindholm@linaro.org> wrote:
>

> On Wed, Jan 23, 2019 at 04:55:28PM +0100, Ard Biesheuvel wrote:

> > On Wed, 23 Jan 2019 at 16:46, Leif Lindholm <leif.lindholm@linaro.org> wrote:

> > >

> > > On Mon, Jan 07, 2019 at 08:15:01AM +0100, Ard Biesheuvel wrote:

> > > > 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/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);

> > >

> >

> > /me pages in the data ...

> >

> > > OK, this whole patch took a few times around the loop before I think I

> > > caught on what was happening.

> > >

> > > I think I'm down to the only things confusing me being:

> > > - The name "Address" to refer to something that is always the start

> > >   address of a 4KB-aligned translation region.

> > >   Is this because the function will be usable in other contexts in

> > >   later patches?

> >

> > I could change it to VirtualAddress if you prefer.

> > It is only passed

> > for TLB maintenance which is only needed at page granularity, and the

> > low bits are shifted out anyway.

>

> Yeah, exactly. It would just be nice if the name reflected that. Not

> sure VirtualAddress does. VirtualBase? PageBase?

>


Well, the argument passed in is called RegionStart, so let's just
stick with that.

> > > - Why drop the & TT_ADDRESS_MASK_DESCRIPTION_TABLE bit here?

> > >   Was it just always pointless and you decided to drop it while you

> > >   were at it?

> >

> > IIRC yes. It is a newly allocated page, so the masking does not do anything.

>

> Yeah, that's fair enough.

> Just made me wonder if anything unobvious was going on :)

>

> /

>     Leif

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
Leif Lindholm Jan. 23, 2019, 4:20 p.m. UTC | #5
On Wed, Jan 23, 2019 at 05:16:57PM +0100, Ard Biesheuvel wrote:
> On Wed, 23 Jan 2019 at 17:13, Leif Lindholm <leif.lindholm@linaro.org> wrote:

> >

> > On Wed, Jan 23, 2019 at 04:55:28PM +0100, Ard Biesheuvel wrote:

> > > On Wed, 23 Jan 2019 at 16:46, Leif Lindholm <leif.lindholm@linaro.org> wrote:

> > > >

> > > > On Mon, Jan 07, 2019 at 08:15:01AM +0100, Ard Biesheuvel wrote:

> > > > > 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/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);

> > > >

> > >

> > > /me pages in the data ...

> > >

> > > > OK, this whole patch took a few times around the loop before I think I

> > > > caught on what was happening.

> > > >

> > > > I think I'm down to the only things confusing me being:

> > > > - The name "Address" to refer to something that is always the start

> > > >   address of a 4KB-aligned translation region.

> > > >   Is this because the function will be usable in other contexts in

> > > >   later patches?

> > >

> > > I could change it to VirtualAddress if you prefer.

> > > It is only passed

> > > for TLB maintenance which is only needed at page granularity, and the

> > > low bits are shifted out anyway.

> >

> > Yeah, exactly. It would just be nice if the name reflected that. Not

> > sure VirtualAddress does. VirtualBase? PageBase?

> >

> 

> Well, the argument passed in is called RegionStart, so let's just

> stick with that.


Sure. With that:
Reviewed-by: Leif Lindholm <leif.lindholm@linaro.org>


> > > > - Why drop the & TT_ADDRESS_MASK_DESCRIPTION_TABLE bit here?

> > > >   Was it just always pointless and you decided to drop it while you

> > > >   were at it?

> > >

> > > IIRC yes. It is a newly allocated page, so the masking does not do anything.

> >

> > Yeah, that's fair enough.

> > Just made me wonder if anything unobvious was going on :)

> >

> > /

> >     Leif

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
Ard Biesheuvel Jan. 28, 2019, 12:29 p.m. UTC | #6
On Wed, 23 Jan 2019 at 17:20, Leif Lindholm <leif.lindholm@linaro.org> wrote:
>

> On Wed, Jan 23, 2019 at 05:16:57PM +0100, Ard Biesheuvel wrote:

> > On Wed, 23 Jan 2019 at 17:13, Leif Lindholm <leif.lindholm@linaro.org> wrote:

> > >

> > > On Wed, Jan 23, 2019 at 04:55:28PM +0100, Ard Biesheuvel wrote:

> > > > On Wed, 23 Jan 2019 at 16:46, Leif Lindholm <leif.lindholm@linaro.org> wrote:

> > > > >

> > > > > On Mon, Jan 07, 2019 at 08:15:01AM +0100, Ard Biesheuvel wrote:

> > > > > > 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/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);

> > > > >

> > > >

> > > > /me pages in the data ...

> > > >

> > > > > OK, this whole patch took a few times around the loop before I think I

> > > > > caught on what was happening.

> > > > >

> > > > > I think I'm down to the only things confusing me being:

> > > > > - The name "Address" to refer to something that is always the start

> > > > >   address of a 4KB-aligned translation region.

> > > > >   Is this because the function will be usable in other contexts in

> > > > >   later patches?

> > > >

> > > > I could change it to VirtualAddress if you prefer.

> > > > It is only passed

> > > > for TLB maintenance which is only needed at page granularity, and the

> > > > low bits are shifted out anyway.

> > >

> > > Yeah, exactly. It would just be nice if the name reflected that. Not

> > > sure VirtualAddress does. VirtualBase? PageBase?

> > >

> >

> > Well, the argument passed in is called RegionStart, so let's just

> > stick with that.

>

> Sure. With that:

> Reviewed-by: Leif Lindholm <leif.lindholm@linaro.org>

>


Thanks

GIven the discussion in the other thread regarding shareability
upgrades of barriers and TLB maintenance instructions when running
under virt, mind if I fold in the hunk below? (and add a mention to
the commit log)

--- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S
+++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S
@@ -35,7 +35,7 @@
   // flush translations for the target address from the TLBs
   lsr   x2, x2, #12
   tlbi  vae\el, x2
-  dsb   sy
+  dsb   nsh

   // re-enable the MMU
   msr   sctlr_el\el, x8
@@ -58,7 +58,7 @@ ASM_FUNC(ArmReplaceLiveTranslationEntry)
   // clean and invalidate first so that we don't clobber
   // adjacent entries that are dirty in the caches
   dc    civac, x0
-  dsb   ish
+  dsb   nsh

   EL1_OR_EL2_OR_EL3(x3)
 1:__replace_entry 1

The first one reduces the scope of the tlb maintenance of a live entry
to the current CPU (and when running under virt, the hypervisor will
upgrade this to inner shareable)
The second one prevents the cache maintenance from being broadcast
unnecessarily, and will be upgraded back to ish in the same way when
running under virt.
_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
Leif Lindholm Jan. 28, 2019, 6:01 p.m. UTC | #7
On Mon, Jan 28, 2019 at 01:29:54PM +0100, Ard Biesheuvel wrote:
> > > > > > > @@ -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);

> > > > > >

> > > > >

> > > > > /me pages in the data ...

> > > > >

> > > > > > OK, this whole patch took a few times around the loop before I think I

> > > > > > caught on what was happening.

> > > > > >

> > > > > > I think I'm down to the only things confusing me being:

> > > > > > - The name "Address" to refer to something that is always the start

> > > > > >   address of a 4KB-aligned translation region.

> > > > > >   Is this because the function will be usable in other contexts in

> > > > > >   later patches?

> > > > >

> > > > > I could change it to VirtualAddress if you prefer.

> > > > > It is only passed

> > > > > for TLB maintenance which is only needed at page granularity, and the

> > > > > low bits are shifted out anyway.

> > > >

> > > > Yeah, exactly. It would just be nice if the name reflected that. Not

> > > > sure VirtualAddress does. VirtualBase? PageBase?

> > > >

> > >

> > > Well, the argument passed in is called RegionStart, so let's just

> > > stick with that.

> >

> > Sure. With that:

> > Reviewed-by: Leif Lindholm <leif.lindholm@linaro.org>

> >

> 

> Thanks

> 

> GIven the discussion in the other thread regarding shareability

> upgrades of barriers and TLB maintenance instructions when running

> under virt, mind if I fold in the hunk below? (and add a mention to

> the commit log)

> 

> --- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S

> +++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S

> @@ -35,7 +35,7 @@

>    // flush translations for the target address from the TLBs

>    lsr   x2, x2, #12

>    tlbi  vae\el, x2

> -  dsb   sy

> +  dsb   nsh


So, this one, definitely. MMU is off, shareability is a shady concept
at best, so it's arguably a fix.

>    // re-enable the MMU

>    msr   sctlr_el\el, x8

> @@ -58,7 +58,7 @@ ASM_FUNC(ArmReplaceLiveTranslationEntry)

>    // clean and invalidate first so that we don't clobber

>    // adjacent entries that are dirty in the caches

>    dc    civac, x0

> -  dsb   ish

> +  dsb   nsh


This one I guess is safe because we know we're never going to get
pre-empted or migrated (because interrupts are disabled and we don't
do that sort of thing)?

If that's the rationale:
Reviewed-by: Leif Lindholm <leif.lindholm@linaro.org>


/
    Leif

>    EL1_OR_EL2_OR_EL3(x3)

>  1:__replace_entry 1

> 

> The first one reduces the scope of the tlb maintenance of a live entry

> to the current CPU (and when running under virt, the hypervisor will

> upgrade this to inner shareable)

> The second one prevents the cache maintenance from being broadcast

> unnecessarily, and will be upgraded back to ish in the same way when

> running under virt.

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
Ard Biesheuvel Jan. 29, 2019, 10:32 a.m. UTC | #8
On Mon, 28 Jan 2019 at 19:01, Leif Lindholm <leif.lindholm@linaro.org> wrote:
>

> On Mon, Jan 28, 2019 at 01:29:54PM +0100, Ard Biesheuvel wrote:

> > > > > > > > @@ -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);

> > > > > > >

> > > > > >

> > > > > > /me pages in the data ...

> > > > > >

> > > > > > > OK, this whole patch took a few times around the loop before I think I

> > > > > > > caught on what was happening.

> > > > > > >

> > > > > > > I think I'm down to the only things confusing me being:

> > > > > > > - The name "Address" to refer to something that is always the start

> > > > > > >   address of a 4KB-aligned translation region.

> > > > > > >   Is this because the function will be usable in other contexts in

> > > > > > >   later patches?

> > > > > >

> > > > > > I could change it to VirtualAddress if you prefer.

> > > > > > It is only passed

> > > > > > for TLB maintenance which is only needed at page granularity, and the

> > > > > > low bits are shifted out anyway.

> > > > >

> > > > > Yeah, exactly. It would just be nice if the name reflected that. Not

> > > > > sure VirtualAddress does. VirtualBase? PageBase?

> > > > >

> > > >

> > > > Well, the argument passed in is called RegionStart, so let's just

> > > > stick with that.

> > >

> > > Sure. With that:

> > > Reviewed-by: Leif Lindholm <leif.lindholm@linaro.org>

> > >

> >

> > Thanks

> >

> > GIven the discussion in the other thread regarding shareability

> > upgrades of barriers and TLB maintenance instructions when running

> > under virt, mind if I fold in the hunk below? (and add a mention to

> > the commit log)

> >

> > --- a/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S

> > +++ b/ArmPkg/Library/ArmMmuLib/AArch64/ArmMmuLibReplaceEntry.S

> > @@ -35,7 +35,7 @@

> >    // flush translations for the target address from the TLBs

> >    lsr   x2, x2, #12

> >    tlbi  vae\el, x2

> > -  dsb   sy

> > +  dsb   nsh

>

> So, this one, definitely. MMU is off, shareability is a shady concept

> at best, so it's arguably a fix.

>

> >    // re-enable the MMU

> >    msr   sctlr_el\el, x8

> > @@ -58,7 +58,7 @@ ASM_FUNC(ArmReplaceLiveTranslationEntry)

> >    // clean and invalidate first so that we don't clobber

> >    // adjacent entries that are dirty in the caches

> >    dc    civac, x0

> > -  dsb   ish

> > +  dsb   nsh

>

> This one I guess is safe because we know we're never going to get

> pre-empted or migrated (because interrupts are disabled and we don't

> do that sort of thing)?

>


Well, we can only get pre-empted under virt, in which case HCR_EL2.BSU
will be set so that barriers are upgraded to inner-shareable. In bare
metal cases, we only care about the local CPU since there are no other
masters (nor DMA capable devices) that care about this cache
maintenance having completed since the page tables are used by the CPU
only.

> If that's the rationale:

> Reviewed-by: Leif Lindholm <leif.lindholm@linaro.org>

>


Thanks

Pushed as f34b38fae614..d5788777bcc7
_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
diff mbox series

Patch

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)