[PATCH RFC v2 07/29] mm: asi: Make __get_current_cr3_fast() ASI-aware

Brendan Jackman
2025-01-10T18:41:08Z

From: Junaid Shahid <junaids@google.com>

When ASI is active, __get_current_cr3_fast() adjusts the returned CR3
value accordingly to reflect the actual ASI CR3.

Signed-off-by: Junaid Shahid <junaids@google.com>
Signed-off-by: Brendan Jackman <jackmanb@google.com>
---
arch/x86/mm/tlb.c | 37 +++++++++++++++++++++++++++++++------
1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/arch/x86/mm/tlb.c b/arch/x86/mm/tlb.c
index 2601beed83aef182d88800c09d70e4c5e95e7ed0..b2a13fdab0c6454c1d9d4e3338801f3402da4191 100644
--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -20,6 +20,7 @@
#include <asm/cache.h>
#include <asm/cacheflush.h>
#include <asm/apic.h>
+#include <asm/asi.h>
#include <asm/perf_event.h>

#include "mm_internal.h"
@@ -197,8 +198,8 @@ static inline unsigned long build_cr3_noflush(pgd_t *pgd, u16 asid,
return build_cr3(pgd, asid, lam) | CR3_NOFLUSH;
}

-noinstr unsigned long build_cr3_pcid_noinstr(pgd_t *pgd, u16 pcid,
- unsigned long lam, bool noflush)
+static __always_inline unsigned long build_cr3_pcid(pgd_t *pgd, u16 pcid,
+ unsigned long lam, bool noflush)
{
u64 noflush_bit = 0;

@@ -210,6 +211,12 @@ noinstr unsigned long build_cr3_pcid_noinstr(pgd_t *pgd, u16 pcid,
return __build_cr3(pgd, pcid, lam) | noflush_bit;
}

+noinstr unsigned long build_cr3_pcid_noinstr(pgd_t *pgd, u16 pcid,
+ unsigned long lam, bool noflush)
+{
+ return build_cr3_pcid(pgd, pcid, lam, noflush);
+}
+
/*
* We get here when we do something requiring a TLB invalidation
* but could not go invalidate all of the contexts. We do the
@@ -1133,14 +1140,32 @@ void flush_tlb_kernel_range(unsigned long start, unsigned long end)
*/
noinstr unsigned long __get_current_cr3_fast(void)
{
- unsigned long cr3 =
- build_cr3(this_cpu_read(cpu_tlbstate.loaded_mm)->pgd,
- this_cpu_read(cpu_tlbstate.loaded_mm_asid),
- tlbstate_lam_cr3_mask());
+ unsigned long cr3;
+ pgd_t *pgd;
+ u16 asid = this_cpu_read(cpu_tlbstate.loaded_mm_asid);
+ struct asi *asi = asi_get_current();
+ u16 pcid;
+
+ if (asi) {
+ pgd = asi_pgd(asi);
+ pcid = asi_pcid(asi, asid);
+ } else {
+ pgd = this_cpu_read(cpu_tlbstate.loaded_mm)->pgd;
+ pcid = kern_pcid(asid);
+ }
+
+ cr3 = build_cr3_pcid(pgd, pcid, tlbstate_lam_cr3_mask(), false);

/* For now, be very restrictive about when this can be called. */
VM_WARN_ON(in_nmi() || preemptible());

+ /*
+ * Outside of the ASI critical section, an ASI-restricted CR3 is
+ * unstable because an interrupt (including an inner interrupt, if we're
+ * already in one) could cause a persistent asi_exit.
+ */
+ VM_WARN_ON_ONCE(asi && asi_in_critical_section());
+
VM_BUG_ON(cr3 != __read_cr3());
return cr3;
}

--
2.47.1.613.gc27f4b7a9f-goog