Re: [PATCH v4 10/10] x86/mm: Try to preserve old TLB entries using PCID

From: Peter Zijlstra
Date: Fri Jul 28 2017 - 09:50:00 EST


On Wed, Jul 05, 2017 at 09:10:00AM -0700, Andy Lutomirski wrote:
> On Wed, Jul 5, 2017 at 5:25 AM, Peter Zijlstra <peterz@xxxxxxxxxxxxx> wrote:
> > On Thu, Jun 29, 2017 at 08:53:22AM -0700, Andy Lutomirski wrote:
> >> +static void choose_new_asid(struct mm_struct *next, u64 next_tlb_gen,
> >> + u16 *new_asid, bool *need_flush)
> >> +{
> >> + u16 asid;
> >> +
> >> + if (!static_cpu_has(X86_FEATURE_PCID)) {
> >> + *new_asid = 0;
> >> + *need_flush = true;
> >> + return;
> >> + }
> >> +
> >> + for (asid = 0; asid < TLB_NR_DYN_ASIDS; asid++) {
> >> + if (this_cpu_read(cpu_tlbstate.ctxs[asid].ctx_id) !=
> >> + next->context.ctx_id)
> >> + continue;
> >> +
> >> + *new_asid = asid;
> >> + *need_flush = (this_cpu_read(cpu_tlbstate.ctxs[asid].tlb_gen) <
> >> + next_tlb_gen);
> >> + return;
> >> + }
> >> +
> >> + /*
> >> + * We don't currently own an ASID slot on this CPU.
> >> + * Allocate a slot.
> >> + */
> >> + *new_asid = this_cpu_add_return(cpu_tlbstate.next_asid, 1) - 1;
> >
> > So this basically RR the ASID slots. Have you tried slightly more
> > complex replacement policies like CLOCK ?
>
> No, mainly because I'm lazy and because CLOCK requires scavenging a
> bit. (Which we can certainly do, but it will further complicate the
> code.) It could be worth playing with better replacement algorithms
> as a followup, though.
>
> I've also considered a slight elaboration of RR in which we make sure
> not to reuse the most recent ASID slot, which would guarantee that, if
> we switch from task A to B and back to A, we don't flush on the way
> back to A. (Currently, if B is not in the cache, there's a 1/6 chance
> we'll flush on the way back.)

How's this?

---
arch/x86/include/asm/mmu.h | 2 +-
arch/x86/include/asm/mmu_context.h | 2 +-
arch/x86/include/asm/tlbflush.h | 2 +-
arch/x86/mm/init.c | 2 +-
arch/x86/mm/tlb.c | 47 ++++++++++++++++++++++++++------------
5 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/arch/x86/include/asm/mmu.h b/arch/x86/include/asm/mmu.h
index bb8c597c2248..9f26ea900df0 100644
--- a/arch/x86/include/asm/mmu.h
+++ b/arch/x86/include/asm/mmu.h
@@ -55,7 +55,7 @@ typedef struct {

#define INIT_MM_CONTEXT(mm) \
.context = { \
- .ctx_id = 1, \
+ .ctx_id = 2, \
}

void leave_mm(int cpu);
diff --git a/arch/x86/include/asm/mmu_context.h b/arch/x86/include/asm/mmu_context.h
index d25d9f4abb15..f7866733875d 100644
--- a/arch/x86/include/asm/mmu_context.h
+++ b/arch/x86/include/asm/mmu_context.h
@@ -137,7 +137,7 @@ static inline void enter_lazy_tlb(struct mm_struct *mm, struct task_struct *tsk)
static inline int init_new_context(struct task_struct *tsk,
struct mm_struct *mm)
{
- mm->context.ctx_id = atomic64_inc_return(&last_mm_ctx_id);
+ mm->context.ctx_id = atomic64_add_return(2, &last_mm_ctx_id);
atomic64_set(&mm->context.tlb_gen, 0);

#ifdef CONFIG_X86_INTEL_MEMORY_PROTECTION_KEYS
diff --git a/arch/x86/include/asm/tlbflush.h b/arch/x86/include/asm/tlbflush.h
index d23e61dc0640..43a4af25d78a 100644
--- a/arch/x86/include/asm/tlbflush.h
+++ b/arch/x86/include/asm/tlbflush.h
@@ -102,7 +102,7 @@ struct tlb_state {
*/
struct mm_struct *loaded_mm;
u16 loaded_mm_asid;
- u16 next_asid;
+ u16 last_asid;

/*
* Access to this CR4 shadow and to H/W CR4 is protected by
diff --git a/arch/x86/mm/init.c b/arch/x86/mm/init.c
index 65ae17d45c4a..570979714d49 100644
--- a/arch/x86/mm/init.c
+++ b/arch/x86/mm/init.c
@@ -812,7 +812,7 @@ void __init zone_sizes_init(void)

DEFINE_PER_CPU_SHARED_ALIGNED(struct tlb_state, cpu_tlbstate) = {
.loaded_mm = &init_mm,
- .next_asid = 1,
+ .last_asid = 0,
.cr4 = ~0UL, /* fail hard if we screw up cr4 shadow initialization */
};
EXPORT_SYMBOL_GPL(cpu_tlbstate);
diff --git a/arch/x86/mm/tlb.c b/arch/x86/mm/tlb.c
index ce104b962a17..aacb87f03428 100644
--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -28,12 +28,28 @@
* Implement flush IPI by CALL_FUNCTION_VECTOR, Alex Shi
*/

-atomic64_t last_mm_ctx_id = ATOMIC64_INIT(1);
+atomic64_t last_mm_ctx_id = ATOMIC64_INIT(2);
+
+static inline u64 asid_ctx_id(int asid)
+{
+ return this_cpu_read(cpu_tlbstate.ctxs[asid].ctx_id) & ~1ULL;
+}
+
+static inline void asid_hit(int asid)
+{
+ this_cpu_or(cpu_tlbstate.ctxs[asid].ctx_id, 1);
+}
+
+static inline bool asid_age(int asid)
+{
+ return this_cpu_xchg(cpu_tlbstate.ctxs[asid].ctx_id, asid_ctx_id(asid)) & 1ULL;
+}

static void choose_new_asid(struct mm_struct *next, u64 next_tlb_gen,
u16 *new_asid, bool *need_flush)
{
u16 asid;
+ int i;

if (!static_cpu_has(X86_FEATURE_PCID)) {
*new_asid = 0;
@@ -42,10 +58,11 @@ static void choose_new_asid(struct mm_struct *next, u64 next_tlb_gen,
}

for (asid = 0; asid < TLB_NR_DYN_ASIDS; asid++) {
- if (this_cpu_read(cpu_tlbstate.ctxs[asid].ctx_id) !=
- next->context.ctx_id)
+ if (asid_ctx_id(asid) != next->context.ctx_id)
continue;

+ asid_hit(asid);
+
*new_asid = asid;
*need_flush = (this_cpu_read(cpu_tlbstate.ctxs[asid].tlb_gen) <
next_tlb_gen);
@@ -53,14 +70,17 @@ static void choose_new_asid(struct mm_struct *next, u64 next_tlb_gen,
}

/*
- * We don't currently own an ASID slot on this CPU.
- * Allocate a slot.
+ * CLOCK - each entry has a single 'used' bit. Cycle through the array
+ * clearing this bit until we find an entry that doesn't have it set.
*/
- *new_asid = this_cpu_add_return(cpu_tlbstate.next_asid, 1) - 1;
- if (*new_asid >= TLB_NR_DYN_ASIDS) {
- *new_asid = 0;
- this_cpu_write(cpu_tlbstate.next_asid, 1);
- }
+ i = this_cpu_read(cpu_tlbstate.last_asid);
+ do {
+ if (--i < 0)
+ i = TLB_NR_DYN_ASIDS - 1;
+ } while (!asid_age(i));
+ this_cpu_write(cpu_tlbstate.last_asid, i);
+
+ *new_asid = i;
*need_flush = true;
}

@@ -125,8 +145,7 @@ void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next,
VM_BUG_ON(__read_cr3() != (__sme_pa(real_prev->pgd) | prev_asid));

if (real_prev == next) {
- VM_BUG_ON(this_cpu_read(cpu_tlbstate.ctxs[prev_asid].ctx_id) !=
- next->context.ctx_id);
+ VM_BUG_ON(asid_ctx_id(prev_asid) != next->context.ctx_id);

if (cpumask_test_cpu(cpu, mm_cpumask(next))) {
/*
@@ -239,9 +258,7 @@ static void flush_tlb_func_common(const struct flush_tlb_info *f,

/* This code cannot presently handle being reentered. */
VM_WARN_ON(!irqs_disabled());
-
- VM_WARN_ON(this_cpu_read(cpu_tlbstate.ctxs[loaded_mm_asid].ctx_id) !=
- loaded_mm->context.ctx_id);
+ VM_WARN_ON(asid_ctx_id(loaded_mm_asid) != loaded_mm->context.ctx_id);

if (!cpumask_test_cpu(smp_processor_id(), mm_cpumask(loaded_mm))) {
/*