为什么要刷 TLB？

写在前面，简单说明一些为什么要刷 TLB？能看到这里的小伙伴，大概率是知道什么是 TLB，TLB 就是一块简单的 cache，但是又与 cache 不同，比如一致性。常有 cache coherence（方便描述，后面直接简写成 CC），但是不见有 TLB coherence。最简单来看，支持 CC，我们就理解成硬件会自动同步 cache 中被修改的数据，保证与其他 cache 中的备份数据，以及内存中的数据的一致性。比如 core 0 上的 cache 数据被更新了，硬件会自动失效其他 core 上相同的备份数据。假如 TLB 与 cache 一样，支持 CC，那么如果 core 0 上的 TLB entry 被失效后，硬件会自动失效其他核心上的 TLB entry（这个 CC 看上去，和 Aarch64 上 TLB flush 的硬件实现有点相似）。

回到正题，由于 TLB 并不支持 CC，所以当内存的数据发生变化时，硬件上并不会失效相应的 TLB entry，这些已经过期的 TLB entry，称之为 stale TLB，如果发生读访问，可能会访问到被释放甚至是其他进程的数据，如果发生写操作，轻则进程异常，重则数据丢失。来个严肃的例子，进程A munmap 一段后备是文件的地址（比如表格 A，写全 A），由于某种不知名的原因没有刷 TLB，直接返回，且被切换到其他子线程 B，子线程申请到了刚才虚拟地址写表格 B（写全 B），这个时候可能会出现数据写串的情况，表格 A 里边出现了完全不可能的字母 B...

因此，在修改 PTE 后，且在虚拟地址空间还未再次分配前，尽快刷新 stale TLB，避免 ABA 问题。

最理想的情况，我们清理了一个 PTE，执行一个 TLB flush 指令，硬件会自动取 snoop 相关的 TLB entry 并失效，等下次访问这个 TLB entry 的时候，不得不触发缺页或从内存 PTE 上去填充 TLB。

在实际的实现上，x86 派系和 arm64 派系有很大差异。

当前是怎么刷 TLB？

当前在 Linux 内核中，需要刷 TLB 的地方，主要在批量 （batch）的基础上，能不刷则不刷，尽量做到一种异步、或者可以说是延迟的刷。举几个简单的例子：

• munmap：munmap 会将一段地址空间从进程的虚拟地址空间去掉（VMA 树上移除），另外，由于刷 TLB 的时候，会考虑到后续的几个待刷的 TLB entry 属于一个进程且连续，一般会等需要取消映射的 PTE 都处理完后，统一刷一次 TLB，则就是批量。当然，也有需要立即刷的时候，比如遇到脏页。
• 内存回收：内存回收的路径上，刷 TLB 就比 munmap、madvise 就要复杂很多。首先，待回收的页面它们不一定属于同一个进程，大多数情况下，可能跨了好几个进程，如果搞一个结构，记录每一个 MM，结构上不是很简洁。因此，在回收的路径上，有批量的刷，更多不同的是，我愿称之为 online flush。什么是 online flush？顾名思义，就是仅仅对那些“正在”核上运行的进程判断是否需要刷 TLB，可能会大于刚才回收的 MM 集合，也可能是小于，但是够用，并且结构解耦了。
• 进程切换：上下文切换估计是刷 TLB 的高频客户，代表函数就是：switch_mm_irqs_off()。上下文切换这块，涉及到一些硬件的东西，比如 PCID/ASID。上面的流程在刷 TLB 的时候会偷些懒，core 0 上触发，但是发现 core 1 上正在运行的不是目标进程，core 1 就会延迟刷，比如延迟到进程切换的时候判断是否需要刷。逃肯定是逃不了的。

下面会列举刷 TLB 的场景补充说明一下。

在当前最新的 Linux 内核中，TLB flush 大同小异，套路无外乎批量刷，延迟刷，另外，投点机，在发 IPI 的时候，判断 online 的进程，过滤掉没必要的同步刷。

munmap/madvise

munmap 调用与 madvise 有些小差异，不会释放页表，先看看 munmap 和 madvise 的调用链路：

do_munmap()
  --> do_vmi_align_munmap()
    --> vms_complete_munmap_vmas()
      --> vms_clear_ptes()
        --> unmap_region()
          --> unmap_vmas()
            --> __zap_vma_range()

do_madvise()
  --> madvise_do_behavior()
    --> madvise_vma_behavior()
      --> madvise_dontneed_free()
        --> madvise_dontneed_single_vma()
          --> zap_vma_range_batched()
            --> __zap_vma_range()

最后都找到遍历页表的统一路径：

zap_p4d_range()
  --> zap_pud_range()
    --> zap_pmd_range()
      --> zap_pte_range()

下面是简化后的核心函数 zap_pte_range()

static unsigned long zap_pte_range(struct mmu_gather *tlb,
				struct vm_area_struct *vma, pmd_t *pmd,
				unsigned long addr, unsigned long end,
				struct zap_details *details)
{
	bool force_flush = false, force_break = false;
	struct mm_struct *mm = tlb->mm;
	...

retry:
	tlb_change_page_size(tlb, PAGE_SIZE);
	init_rss_vec(rss);
	start_pte = pte = pte_offset_map_lock(mm, pmd, addr, &ptl);
	if (!pte)
		return addr;

	flush_tlb_batched_pending(mm); /* 标记1 */
	lazy_mmu_mode_enable(); /* 标记2 */
	do {
		...
		nr = do_zap_pte_range(tlb, vma, pte, addr, end, details, rss,
				      &force_flush, &force_break, &any_skipped);
		...
	} while (pte += nr, addr += PAGE_SIZE * nr, addr != end);
	...
	lazy_mmu_mode_disable();

	/* Do the actual TLB flush before dropping ptl */
	if (force_flush) {
		tlb_flush_mmu_tlbonly(tlb); /* 标注3 */
		tlb_flush_rmaps(tlb, vma);
	}
	pte_unmap_unlock(start_pte, ptl);

	/*
	 * If we forced a TLB flush (either due to running out of
	 * batch buffers or because we needed to flush dirty TLB
	 * entries before releasing the ptl), free the batched
	 * memory too. Come back again if we didn't do everything.
	 */
	if (force_flush)
		tlb_flush_mmu(tlb);

    ...

	return addr;
}

简化后的重点知识不多，这里拧出三点：

1. 标记 1：调用flush_tlb_batched_pending 函数刷同时进行的回收流程中的挂起的待刷操作，个人觉的没必要，延迟到 vma 释放到 vma tree 前来刷新即可，也许是个可以改进的地方，减少一些 IPI；
2. 标记 2：lazy_mmu_mode_enable 函数，x86 当前还没怎么支持，可以忽略；个人觉的在内核态，大多数时候不会访问当前进程用户态的页表，在进入内核态后，可以在完全确认的路径上加上 lazy mode，这样在其他核心发 IPI 的时候减少一些不必要的 tlb shootdown 影响。
3. 标记 3：这里有一个强制刷，这种情况仅仅对文件页而言。匿名页估计还好，漏就漏掉了，不需要写回到磁盘，文件页会刷会到磁盘，牵扯出一些存储问题可能就不是搞内存的能兜得住了。

这部分逻辑属于 madvise 和 munmap 共用的，除非遇到脏页，否则都是批量刷。

前不久，madvise 批量刷 TLB 有一些调整，当前 madvise 流程的批量刷可以在 do_madvise 函数中看到，没什么需要补充的。

int do_madvise(struct mm_struct *mm, unsigned long start, size_t len_in, int behavior)
{
	int error;
	struct mmu_gather tlb;
	struct madvise_behavior madv_behavior = {
		.mm = mm,
		.behavior = behavior,
		.tlb = &tlb,
	};

	if (madvise_should_skip(start, len_in, behavior, &error))
		return error;
	error = madvise_lock(&madv_behavior);
	if (error)
		return error;
	madvise_init_tlb(&madv_behavior); /* <-- 这里初始化是否支持批量刷 */
	error = madvise_do_behavior(start, len_in, &madv_behavior);
	madvise_finish_tlb(&madv_behavior); /* <-- 这里调用tlb_finish_mmu */
	madvise_unlock(&madv_behavior);

	return error;
}

回收流程

回收流程与 munmap/madvise 有点不同，最显眼的就是 mm 不同。

回收过程针对某个 memcg，再到具体的 LRU 去回收，是一种反向判断，比如是先获取到页面，然后根据页面判断是否能回收。这种情况下，在回收的页面链表中，跨了好几个 mm。另外，回收过程回收的页面会立即进入到 PCP，然后被再分配，所以回收流程刷 TLB 需要更勤快些。

先说结论，回收流程也是批量刷，但是具体实现上有些小同大异。

shrink_folio_list()
 1--> try_to_unmap_one()
    --> set_tlb_ubc_flush_pending()
      --> inc_mm_tlb_gen(mm) /* 增加当前mm->tlb_gen */
 2--> try_to_unmap_flush()
    --> arch_tlbbatch_flush()
        --> flush_tlb_multi()
            --> native_flush_tlb_multi()

回收流程刷 TLB 的流程大致如上。这块用到了 mm->tlb_gen，在 per cpu 那还有一个该进程的 tlb_gen，这个变量主要有两个作用，这里是一个，另外一个在上下文切换的时候再聊。在回收流程中，暂时没有记录哪些 mm 在这批次待刷的名单中，因此无法判断哪些 mm 需要刷，所以批量触发到 try_to_unmap_flush 函数时，就需要一个策略怎么刷才安全。

答案在native_flush_tlb_multi 函数中。从try_to_unmap_flush 函数中没有传入任何待刷的 mm 信息，仅给了一个TLB_FLUSH_ALL。native_flush_tlb_multi 函数用了一种简单粗暴的方式，凡是 online 的进程，能刷抖刷。能刷需要两个条件，进程正在运行且 mm->tlb_gen 大于其 per cpu 维护的 tlb_gen，不难理解，说明进程在某个流程中需要刷，但是还未来得及刷。这个集合可能大于也可能小于当前回收 mm 集合。另外，回收过程明确标注了TLB_FLUSH_ALL，表示当判断这个 mm 需要刷的时候，刷该进程所有的 tlb entry，不区分是否是 stale。

这块简单夹带点个人私货，不能保真：

1. 回收过程实际刷的 TLB 和前面回收的 mm 不一定相关，但在概率上，该有的都有。另外，对于刚刚回收的 mm，但try_to_unmap_flush 流程并没有刷 TLB，可能是调度出去了，也可能是退出。对于调度出去，该来的迟早要来，在它下次被调度回时，还得刷 TLB。无论下次这个进程在哪个核心被调度，由于 mm->tlb_gen 增加了，怎么也跑不掉。
2. 大多地方把回收这个地方的方式称为 global flush，个人感觉不够准确，其实称为 online flush 或者 all online flush，更能顾名思义。
3. 这里是否需要实现一个待刷的 mm 链表，传入try_to_unmap_flus，有针对性的刷值得讨论。一方面结构需要调整不少，另外这些省下来的 IPI 是否有收益还是值得商榷。另外，各大厂商都有自己的主动回收方案，将这种精确性的 mm 链表实现到主动回收也未尝不能用，这样就能减少上游的修改。

上下文切换

在讨论上下文切换中的 TLB flush 时，我们需要一些先验知识。

1. mm_cpumask：进程中有一个变量记录了该进程正在哪些核心上运行，在调度出去或者调度回时，会更新该变量。
2. ASID/PCID：就以 x86 PCID 为准，硬件上支持 4096 个 PCID，内核在每个核心上仅仅用 6 个。这 6 个 PCID 所有运行在该核心上的进程共享。当一个进程被调度回时，会优先寻找上一轮使用的 ASID，如果已经被其他进程占用，则也会去抢占其他进程的 ASID；
3. cpu_tlbstate.ctxs[asid].tlb_gen：用了 6 个槽位，其中 asid 这个槽位记录的 tlb_gen 值。主要用来判断当进程被调度回时并寻找到上一轮的 ASID，是否需要刷 TLB。

switch_mm_irqs_off 函数太长，这里就拧出choose_new_asid 简单过一遍：

static struct new_asid choose_new_asid(struct mm_struct *next, u64 next_tlb_gen)
{
	struct new_asid ns;
	u16 asid;

	if (!static_cpu_has(X86_FEATURE_PCID)) {
		ns.asid = 0;
		ns.need_flush = 1;
		return ns;
	}

	/*
	 * TLB consistency for global ASIDs is maintained with hardware assisted
	 * remote TLB flushing. Global ASIDs are always up to date.
	 */
	if (cpu_feature_enabled(X86_FEATURE_INVLPGB)) {
		u16 global_asid = mm_global_asid(next);

		if (global_asid) {
			ns.asid = global_asid;
			ns.need_flush = 0;
			return ns;
		}
	}

	if (this_cpu_read(cpu_tlbstate.invalidate_other))
		clear_asid_other();

	for (asid = 0; asid < TLB_NR_DYN_ASIDS; asid++) { /* 标注1：寻找上一轮 ASID */
		if (this_cpu_read(cpu_tlbstate.ctxs[asid].ctx_id) !=
		    next->context.ctx_id)
			continue;

		ns.asid = asid; /* 标注2：判断是否需要刷 TLB */
		ns.need_flush = (this_cpu_read(cpu_tlbstate.ctxs[asid].tlb_gen) < next_tlb_gen);
		return ns;
	}

	/*
	 * We don't currently own an ASID slot on this CPU.
	 * Allocate a slot.
	 */
	ns.asid = this_cpu_add_return(cpu_tlbstate.next_asid, 1) - 1;
	if (ns.asid >= TLB_NR_DYN_ASIDS) { /* 标注3：抢占其他的 ASID */
		ns.asid = 0;
		this_cpu_write(cpu_tlbstate.next_asid, 1);
	}
	ns.need_flush = true;

	return ns;
}

这里主要有三点：

• 标注 1：进程被调度回时，会先遍历 6 个 ASID 槽位，能复用则复用；
• 标注 2：哪怕复用了上一轮的 ASID，保不齐，进程不在家的这段时间，其他核心上的子线程或者回收线程修改了页表，所以这里需要兜底判断是否需要刷 TLB（在写 CR3 的时候，通过 NOFLUSH 标注刷还是不刷）。
• 标注 3：没有找到上一轮的信息，只能抢占其他进程的 ASID。这里没有什么策略，挨个来。这种情况必须刷，必须防止吃席吃到了上一轮的残羹剩饭。

上下文切换过程中的 tlb_gen 判断，其实就是为之前 online flush 的投机行为兜底。这里 Aarch64 和 x86 不太一样，Aarch64 不需要软件去判断 mm cpumask，硬件会自动判断哪些该刷或不刷。这种方式其实也并非全好，在核多的时候，大量的 tlb snoop 消息也会有不小的开销。

画一张简图总结上面讨论到的 tlb flush 情况。这些也是 Linux 内核中比较明显的机制。其中一些细节，比如在上下文切换中，刷与不刷也会考虑到上次运行的是内核线程或其他。主要就一个原则，新产生的 stale TLBs 是否有可能被访问到。

写在最后

x86 给了软件更多的自由，将是否发 IPI 这种情况交给了软件来执行。看上去是一场开卷考试，但是标准答案还只有一个，为了安全，你别无选择，只能都刷。这种自由，只有在拥有更多的情况下才能发挥作用。比如，提供一个指令能获取当前哪些核有该页表项的缓存，甚至更理想情况下，core A 修改了页表，硬件能自动广播到需要广播的核心上，类似 invalidate queue，其他核心访问到该 stale TLB 时，在 L1 TLB 匹配时检查 invalidate queue，以此可以更大程度简化软件逻辑。