Linux Kernel HMM 深度工程实战：CPU-GPU 统一虚拟地址空间从原理到驱动实现

在大规模 GPU 计算场景中，传统的「CPU 内存 设备显存」二分法已成为性能瓶颈的主要推手。每次 kernel launch 都依赖 cudaMemcpy 进行显式数据搬运，不仅编程模型笨拙，更因 PCIe 2.0/3.0 的有限带宽难以匹配计算吞吐。Linux 内核的 HMM（Heterogeneous Memory Management）子系统试图从内核层面彻底解决这一问题：让 CPU 和 GPU 共享同一套页表，按需迁移、按需回退，实现"零拷贝、按需分页"的统一虚拟地址空间。本文将深入拆解 HMM 的内核实现机制、mmu_notifier 回调体系、页面迁移协议，并结合 NVIDIA Pascal 以上架构的实际案例写一份最小可工作的设备驱动骨架代码。

一、为什么需要 HMM —— 一个由带宽倒逼出的架构变革

1.1 统一虚拟内存（UVM）的商业动机

NVIDIA 在 2014 年随 Maxwell 架构引入了 UVM（Unified Virtual Addressing），但其底层实现完全私有：驱动自行维护 RADIX 树、自行处理 page fault、自行发起 PCIe DMA。这带来三个严重问题：

1. 内核页表感知缺失：自行维护的地址空间与 CPU 真实页表不同步，导致 get_user_pages 等内核接口完全无法作用于设备映射的地址；
2. mmu_notifier 不参与：进程 fork() 后的 COW、KSM 合并、KSM 拆分、memory compaction 等行为时，设备驱动完全无法响应；
3. 热插拔困难：GPU 被热拔或 FLR（Function Level Reset）后，整个地址黑洞需要驱动暴力清理，缺乏与 mm_struct 的标准协作接口。

1.2 HMM 的设计目标

HMM 于 Linux 4.14 引入后经历 5.x 重大重构，其核心目标可以浓缩为：

让异构设备成为 struct mm_struct 的一等公民，与普通 CPU 进程参与相同的 VMA（Virtual Memory Area）生命周期，通过 mmu_notifier 自动接收地址空间事件通知。

达成这一目标后，get_user_pages(page, FOLL_GET) 现在可以直接返回 device-mapped page、fork() 自动触发设备端 COW、甚至 madvise(MADV_DONTNEED) 都会通知驱动释放设备端映射。

二、HMM 架构全景

2.1 核心数据结构

┌──────────────────────────────────────────────────────┐
│                  进程 mm_struct                       │
│  ┌─────────────────────────────────────────────────┐ │
│  │             hmm (per-process)                    │ │
│  │  hmm_ranges: RB tree of [start, end)            │ │
│  │  hmm_mirrors: list of attached devices           │ │
│  └─────────────────────────────────────────────────┘ │
└──────────────────────────────────────────────────────┘
                        │
                        ▼
┌──────────────────────────────────────────────────────┐
│   Device A          Device B           Device C      │
│  hmm_mirror       hmm_mirror        hmm_mirror       │
│  ┌───────────┐   ┌───────────┐    ┌───────────┐      │
│  │ dev-pte   │   │ dev-pte   │    │ dev-pte   │      │
│  │ 镜像页表   │   │ 镜像页表   │    │ 镜像页表   │      │
│  └───────────┘   └───────────┘    └───────────┘      │
└──────────────────────────────────────────────────────┘

struct hmm（include/linux/hmm.h）：

struct hmm {
    struct mm_struct *mm;       /* 绑定的进程地址空间 */
    struct kref kref;           /* 引用计数，进程释放时销毁 */
    struct hmm_range ranges;    /* 当前有效的地址范围树 */
    struct hmm_cpu_deferred     /* CPU fault 延迟处理（bottom-half） */
};

struct hmm_mirror：每个接入 HMM 的设备实例维护自己的镜像对象，内含 struct hmm_ops *ops。ops 集合是驱动必须实现的最小语义接口：

struct hmm_ops {
    /* CPU 访问了一个 device-only 的页时，要求驱动把数据迁移回 RAM */
    void (*migrate_to_ram)(struct hmm_mirror *mirror,
                          struct hmm_range *range);

    /* 将指定范围的页从 RAM 迁移到设备本地存储 */
    int (*migrate_to_device)(struct hmm_mirror *mirror,
                            struct hmm_range *range);

    /* 释放设备端对指定地址范围的 PTE 引用 */
    void (*release)(struct hmm_mirror *mirror,
                   struct hmm_range *range);

    /* 设备侧发生缺页中断时，驱动需要返回 device page */
    int (*device_fault_entry)(struct hmm_mirror *mirror,
                              struct hmm_range *range);
};

struct hmm_range：一次迁移/故障操作的工作集，包含 [start, end) 范围、一个 struct page ** 数组以及 fault 标志位。

2.2 hmm_range_fault() —— 核心入口

hmm_range_fault() 是驱动在四种场景下都必须调用的统一入口：

/**
 * hmm_range_fault - 启动一次地址范围操作
 * @range: 当前操作上下文
 *
 * 返回 0 表示所有页可访问；-EAGAIN 表示页面需要迁移；
 * -EBUSY 表示页面正在被另一个操作锁定。
 */
int hmm_range_fault(struct hmm_range *range);

内部工作原理：

1. mmap_read_lock(mm) 获取 mm 读锁；
2. 通过 find_vma() 和 follow_page() 遍历当前 VMA；
3. 若页表项为 migration entry，调用 migrate_vma_collect_pagenotify() 定位器；
4. 若页在 device memory，回调 mirror-