[RFC PATCH v3 3/7] mm/asi: Improve TLB flushing when switching to an ASI pagetable

[Alexandre Chartre]

When switching to an ASI pagetable, the TLB doesn't need to be flushed
if it was previously used with the same PCID. So, to avoid unnecessary
TLB flushing, we track which pagetables are used with the different
ASI PCIDs. If an ASI PCID is being used with a different ASI pagetable,
or if we have a new generation of the same ASI pagetable, then the TLB
needs to be flushed. This behavior is similar to the context tracking
done when switching mm.

Signed-off-by: Alexandre Chartre <alexandre.chartre@xxxxxxxxxx>
---
 arch/x86/include/asm/asi.h |   23 +++++++++++++++++++++++
 arch/x86/mm/asi.c          |   34 ++++++++++++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/arch/x86/include/asm/asi.h b/arch/x86/include/asm/asi.h
index 29b745a..bcfb68e 100644
--- a/arch/x86/include/asm/asi.h
+++ b/arch/x86/include/asm/asi.h
@@ -46,8 +46,26 @@
 
 #include <asm/asi_session.h>
 
+/*
+ * ASI_NR_DYN_ASIDS is the same as TLB_NR_DYN_ASIDS. We can't directly
+ * use TLB_NR_DYN_ASIDS because asi.h and tlbflush.h can't both include
+ * each other.
+ */
+#define ASI_TLB_NR_DYN_ASIDS	6
+
+struct asi_tlb_pgtable {
+	u64 id;
+	u64 gen;
+};
+
+struct asi_tlb_state {
+	struct asi_tlb_pgtable	tlb_pgtables[ASI_TLB_NR_DYN_ASIDS];
+};
+
 struct asi_type {
 	int			pcid_prefix;	/* PCID prefix */
+	struct asi_tlb_state	*tlb_state;	/* percpu ASI TLB state */
+	atomic64_t		last_pgtable_id; /* last id for this type */
 };
 
 /*
@@ -58,8 +76,11 @@ struct asi_type {
  * specified type.
  */
 #define DEFINE_ASI_TYPE(name, pcid_prefix)			\
+	DEFINE_PER_CPU(struct asi_tlb_state, asi_tlb_ ## name);	\
 	struct asi_type asi_type_ ## name = {			\
 		pcid_prefix,					\
+		&asi_tlb_ ## name,				\
+		ATOMIC64_INIT(1),				\
 	};							\
 	EXPORT_SYMBOL(asi_type_ ## name)
 
@@ -76,6 +97,8 @@ struct asi_type {
 struct asi {
 	struct asi_type		*type;		/* ASI type */
 	pgd_t			*pagetable;	/* ASI pagetable */
+	u64			pgtable_id;	/* ASI pagetable ID */
+	atomic64_t		pgtable_gen;	/* ASI pagetable generation */
 	unsigned long		base_cr3;	/* base ASI CR3 */
 };
 
diff --git a/arch/x86/mm/asi.c b/arch/x86/mm/asi.c
index 9fbc921..cf0d122 100644
--- a/arch/x86/mm/asi.c
+++ b/arch/x86/mm/asi.c
@@ -25,6 +25,8 @@ struct asi *asi_create(struct asi_type *type)
 		return NULL;
 
 	asi->type = type;
+	asi->pgtable_id = atomic64_inc_return(&type->last_pgtable_id);
+	atomic64_set(&asi->pgtable_gen, 0);
 
 	return asi;
 }
@@ -61,6 +63,33 @@ void asi_set_pagetable(struct asi *asi, pgd_t *pagetable)
 }
 EXPORT_SYMBOL(asi_set_pagetable);
 
+/*
+ * Update ASI TLB flush information for the specified ASI CR3 value.
+ * Return an updated ASI CR3 value which specified if TLB needs to
+ * be flushed or not.
+ */
+static unsigned long asi_update_flush(struct asi *asi, unsigned long asi_cr3)
+{
+	struct asi_tlb_pgtable *tlb_pgtable;
+	struct asi_tlb_state *tlb_state;
+	s64 pgtable_gen;
+	u16 pcid;
+
+	pcid = asi_cr3 & ASI_KERNEL_PCID_MASK;
+	tlb_state = get_cpu_ptr(asi->type->tlb_state);
+	tlb_pgtable = &tlb_state->tlb_pgtables[pcid - 1];
+	pgtable_gen = atomic64_read(&asi->pgtable_gen);
+	if (tlb_pgtable->id == asi->pgtable_id &&
+	    tlb_pgtable->gen == pgtable_gen) {
+		asi_cr3 |= X86_CR3_PCID_NOFLUSH;
+	} else {
+		tlb_pgtable->id = asi->pgtable_id;
+		tlb_pgtable->gen = pgtable_gen;
+	}
+
+	return asi_cr3;
+}
+
 static void asi_switch_to_asi_cr3(struct asi *asi)
 {
 	unsigned long original_cr3, asi_cr3;
@@ -72,10 +101,11 @@ static void asi_switch_to_asi_cr3(struct asi *asi)
 	original_cr3 = __get_current_cr3_fast();
 
 	/* build the ASI cr3 value */
-	asi_cr3 = asi->base_cr3;
 	if (boot_cpu_has(X86_FEATURE_PCID)) {
 		pcid = original_cr3 & ASI_KERNEL_PCID_MASK;
-		asi_cr3 |= pcid;
+		asi_cr3 = asi_update_flush(asi, asi->base_cr3 | pcid);
+	} else {
+		asi_cr3 = asi->base_cr3;
 	}
 
 	/* get the ASI session ready for entering ASI */
-- 
1.7.1