[PATCH] arm64: mm: Optimize querying asid from reserved_asids

@@ -101,7 +101,7 @@ static void set_reserved_asid_bits(void)
 #define asid_gen_match(asid) \
 	(!(((asid) ^ atomic64_read(&asid_generation)) >> asid_bits))
 
-static void flush_context(void)
+static void flush_context(u64 generation)
 {
 	int i;
 	u64 asid;

@@ -120,6 +120,8 @@ static void flush_context(void)
 		 */
 		if (asid == 0)
 			asid = per_cpu(reserved_asids, i);
+
+		asid = generation | (asid & ~ASID_MASK);
 		__set_bit(ctxid2asid(asid), asid_map);
 		per_cpu(reserved_asids, i) = asid;
 	}

@@ -131,24 +133,21 @@ static void flush_context(void)
 	cpumask_setall(&tlb_flush_pending);
 }
 
-static bool check_update_reserved_asid(u64 asid, u64 newasid)
+static bool check_update_reserved_asid(u64 newasid)
 {
-	int cpu;
+	int cpu, cur_cpu = smp_processor_id();
 	bool hit = false;
 
-	/*
-	 * Iterate over the set of reserved ASIDs looking for a match.
-	 * If we find one, then we can update our mm to use newasid
-	 * (i.e. the same ASID in the current generation) but we can't
-	 * exit the loop early, since we need to ensure that all copies
-	 * of the old ASID are updated to reflect the mm. Failure to do
-	 * so could result in us missing the reserved ASID in a future
-	 * generation.
-	 */
+	if (per_cpu(reserved_asids, cur_cpu) == newasid)
+		return true;
+
 	for_each_possible_cpu(cpu) {
-		if (per_cpu(reserved_asids, cpu) == asid) {
+		if (cpu == cur_cpu)
+			continue;
+
+		if (per_cpu(reserved_asids, cpu) == newasid) {
 			hit = true;
-			per_cpu(reserved_asids, cpu) = newasid;
+			break;
 		}
 	}
 

@@ -168,7 +167,7 @@ static u64 new_context(struct mm_struct *mm)
 		 * If our current ASID was active during a rollover, we
 		 * can continue to use it and this was just a false alarm.
 		 */
-		if (check_update_reserved_asid(asid, newasid))
+		if (check_update_reserved_asid(newasid))
 			return newasid;
 
 		/*

@@ -201,7 +200,7 @@ static u64 new_context(struct mm_struct *mm)
 	/* We're out of ASIDs, so increment the global generation count */
 	generation = atomic64_add_return_relaxed(ASID_FIRST_VERSION,
 						 &asid_generation);
-	flush_context();
+	flush_context(generation);
 
 	/* We have more ASIDs than CPUs, so this will always succeed */
 	asid = find_next_zero_bit(asid_map, NUM_USER_ASIDS, 1);