第4章:NPU内存管理:DMA缓冲区分配与管理、IOMMU/SMMU地址映射、连续内存分配器(CMA)、用户态与内核态内存共享

各位同学,咱们今天聊点硬核的——NPU的内存管理。

说实话,我刚开始接触NPU驱动时,最头疼的就是内存这块。CPU那边跑得好好的程序,一到NPU这边就崩,查来查去,十有八九是内存没管好。你想想看,NPU要处理的数据量动不动就是几十上百兆,甚至上G,这可不是闹着玩的。

4.1 DMA缓冲区分配与管理

先说说DMA缓冲区。NPU和CPU之间传数据,靠的就是DMA。但这里有个坑——DMA需要物理连续的内存。

为什么?因为DMA控制器不像CPU有MMU,它只能按物理地址连续地搬数据。你给它一个不连续的缓冲区,它直接就懵了。

我在项目中遇到过这么个事:某次做视频分析,NPU需要从摄像头DMA一帧1080p的图像。我一开始用kmalloc分配缓冲区,结果系统跑了一段时间后,kmalloc分配不出大块连续内存了。嗯,这就是碎片化问题。

核心要点: DMA缓冲区分配,首选dma_alloc_coherent(),它保证物理连续且cache一致性处理好。

代码示例:

// 分配DMA缓冲区
dma_addr_t dma_handle;
void *cpu_addr;

cpu_addr = dma_alloc_coherent(dev, size, &dma_handle, GFP_KERNEL);
if (!cpu_addr) {
    dev_err(dev, "DMA buffer allocation failed\n");
    return -ENOMEM;
}

// 使用完后释放
dma_free_coherent(dev, size, cpu_addr, dma_handle);

这里有个细节:dma_alloc_coherent返回的CPU地址和DMA地址是两回事。CPU用cpu_addr访问,DMA控制器用dma_handle访问。千万别搞混了,我曾经见过有人把dma_handle当CPU指针用,结果系统直接挂了。

4.2 IOMMU/SMMU地址映射

好,现在问题来了。如果NPU需要访问不连续的物理内存怎么办?或者,你想让NPU访问用户态进程的内存?这时候就需要IOMMU(也叫SMMU)出场了。

说白了,IOMMU就是给外设用的MMU。它能把不连续的物理内存映射成连续的IO虚拟地址,还能做权限控制。

我个人的习惯是:只要NPU支持SMMU,就一定要打开。为什么?因为有了SMMU,你就不用纠结物理连续了,系统可以给你分配任意物理页,SMMU帮你搞定映射。

经验之谈: 配置SMMU时,页表大小和CPU的页表大小保持一致(通常是4KB),这样能减少TLB miss。

SMMU映射的核心流程:

  1. 分配IOVA(IO虚拟地址)区间
  2. 获取物理页的PFN(页帧号)
  3. 调用iommu_map()建立映射
  4. 使用完毕后调用iommu_unmap()解除映射
// SMMU映射示例
struct iommu_domain *domain;
unsigned long iova;
phys_addr_t phys;
size_t size;

domain = iommu_domain_alloc(bus);
if (!domain) {
    pr_err("Failed to allocate IOMMU domain\n");
    return -ENOMEM;
}

// 分配IOVA区间
iova = iommu_dma_alloc_iova(domain, size, dev->coherent_dma_mask);
if (!iova) {
    pr_err("Failed to allocate IOVA\n");
    goto err_free_domain;
}

// 建立映射
ret = iommu_map(domain, iova, phys, size, IOMMU_READ | IOMMU_WRITE);
if (ret) {
    pr_err("IOMMU map failed: %d\n", ret);
    goto err_free_iova;
}

这里要注意:SMMU的页表刷新是有开销的。我在做实时推理时,就踩过这个坑——频繁map/unmap导致性能下降。后来改成预分配IOVA池,性能才上去。

4.3 连续内存分配器(CMA)

说到连续内存,就不得不提CMA。这是内核提供的一个机制:预留一大块内存,平时给系统用,但需要时可以回收给DMA用。

你想想看,如果系统里每个驱动都自己预留DMA缓冲区,那内存浪费得多严重。CMA就是解决这个问题的——它让内存可以被复用。

避坑指南: CMA区域的大小要合理设置。太小了不够用,太大了又浪费。我一般按NPU最大工作负载的1.5倍来设。

CMA的使用方式:

// 设备树中配置CMA
reserved-memory {
    #address-cells = <2>;
    #size-cells = <2>;
    
    npu_cma: npu_cma_region {
        compatible = "shared-dma-pool";
        reusable;
        size = <0 0x10000000>;  // 256MB
        alignment = <0x200000>;
        linux,cma-default;
    };
};

// 驱动中获取CMA区域
struct device *dev;
struct cma *cma;

cma = dev_get_cma_area(dev);
if (!cma) {
    dev_err(dev, "No CMA area available\n");
    return -ENODEV;
}

// 从CMA分配内存
unsigned long pfn = cma_alloc(cma, count, 0, GFP_KERNEL);
if (!pfn) {
    dev_err(dev, "CMA allocation failed\n");
    return -ENOMEM;
}

我曾经在一个项目中,CMA设了512MB,结果系统启动时其他模块报内存不足。后来发现是CMA区域和内核线性映射区冲突了。解决办法是调整CMA的物理地址范围,避开内核关键区域。

4.4 用户态与内核态内存共享

最后这个部分,我觉得是最有意思的。NPU驱动最终是要给应用程序用的,那用户态怎么把数据传给NPU?

传统做法是:用户态 -> 内核态(copy_from_user)-> DMA缓冲区 -> NPU。但这样数据要拷贝两次,性能损失很大。

更好的做法是:让用户态和NPU共享同一块物理内存。怎么做?用mmap。

我个人的经验是:先在内核中分配DMA缓冲区,然后通过mmap映射到用户态。这样用户态程序可以直接操作这块内存,NPU也能直接访问,零拷贝。

// 内核驱动的mmap实现
static int npu_mmap(struct file *filp, struct vm_area_struct *vma)
{
    struct npu_device *npu = filp->private_data;
    unsigned long size = vma->vm_end - vma->vm_start;
    unsigned long pfn = virt_to_phys(npu->buffer) >> PAGE_SHIFT;
    
    // 设置内存属性为uncached,保证一致性
    vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
    
    // 映射物理内存到用户空间
    if (remap_pfn_range(vma, vma->vm_start, pfn, 
                        size, vma->vm_page_prot)) {
        dev_err(npu->dev, "mmap failed\n");
        return -EAGAIN;
    }
    
    return 0;
}
关键点: 共享内存一定要处理好cache一致性。NPU和CPU对同一块内存的访问,如果cache策略不一致,数据就乱了。我建议用dma_alloc_coherent或设置pgprot_noncached。

用户态的使用方式:

// 用户态程序
int fd = open("/dev/npu", O_RDWR);
void *buf = mmap(NULL, size, PROT_READ | PROT_WRITE, 
                 MAP_SHARED, fd, 0);

// 直接往buf里写数据
memcpy(buf, input_data, input_size);

// 通知NPU处理
ioctl(fd, NPU_RUN, &config);

// 处理完成后,结果就在buf里
memcpy(output_data, buf, output_size);

munmap(buf, size);
close(fd);

这里有个坑:用户态mmap后,如果进程崩溃了,内核要记得在release时释放映射。不然就会内存泄漏。我早期做驱动时就漏过,结果系统跑了几天后OOM了。

知识体系总览

下面这张图,是我梳理的NPU内存管理核心逻辑。你看一眼,基本就明白这几块内容的关系了。

NPU内存管理知识体系 用户态应用程序 NPU内核驱动(内存管理核心) DMA缓冲区管理 dma_alloc_coherent 物理连续 / cache一致性 IOMMU/SMMU地址映射 iommu_map / iova分配 虚拟地址 / 权限控制 CMA连续内存分配器 cma_alloc / 内存复用 预留区域 / 动态回收 用户态/内核态共享 mmap / 零拷贝 cache一致性 / 安全控制 NPU硬件 + 系统内存 + SMMU硬件

这张图把整个流程串起来了:用户态通过mmap和内核共享内存,内核用DMA缓冲区或CMA分配物理内存,SMMU负责地址映射,最终NPU硬件直接访问数据。

好了,这一章的内容就到这。内存管理是NPU驱动的基石,这块搞不定,后面推理加速什么的都是空中楼阁。建议你把代码示例都跑一遍,遇到问题随时交流。

想要更多资料,联系可可