开发者看到的是torch.randn(1024, 1024, devicenpu)驱动做的事PCIe 枚举找到 NPU 设备 → 初始化 HBM 页表 → 分配物理内存 → 创建 Stream 上下文 → 提交 DMA 命令到硬件队列 → 返回用户态句柄。每一步出错都会变成torch.npu的 RuntimeError 或更糟——静默性能退化。PCIe 设备枚举从 Vendor ID 到 Device Handle// driver/src/pcie/pcie_enum.cpp// PCIe 枚举扫描所有 PCIe 总线找到 Ascend NPU 设备// Ascend NPU 的 Vendor ID 0x19E5华为/昇腾Device ID 根据型号不同structAscendDevice{uint16_tvendor_id;// 0x19E5uint16_tdevice_id;// 0xD100 Ascend 910, 0xD110 910B, 0xD200 950uint8_tbus;// PCIe 总线号uint8_tdevice;// 设备号uint8_tfunction;// 功能号uint64_tbar0_base;// BAR0 基地址寄存器映射通常 32KBuint64_tbar0_size;uint64_tbar2_base;// BAR2 基地址HBM 映射通常 64GBuint64_tbar2_size;intnuma_node;// NUMA 节点PCIe 拓扑决定intcore_count;// AI Core 数量910 32, 910B 64uint64_thbm_size;// HBM 总大小910 56GB};// 枚举所有 Ascend NPU 设备std::vectorAscendDeviceEnumerateAscendDevices(){std::vectorAscendDevicedevices;intdevice_index0;// 扫描所有 PCIe 总线bus 0..255, device 0..31, function 0..7for(intbus0;bus256;bus){for(intdev0;dev32;dev){for(intfunc0;func8;func){// 读取 PCIe 配置空间的 Vendor ID Device IDuint32_tvendor_deviceReadPCIConfig(bus,dev,func,0x00);uint16_tvendor_idvendor_device0xFFFF;uint16_tdevice_idvendor_device16;if(vendor_id!0x19E5)continue;// 不是昇腾设备AscendDevice d;d.vendor_idvendor_id;d.device_iddevice_id;d.busbus;d.devicedev;d.functionfunc;d.device_indexdevice_index;// 读取 BARBase Address Register// BAR0offset 0x10 in PCIe config spaceuint64_tbar0ReadPCIConfig64(bus,dev,func,0x10);d.bar0_basebar0~0xF;// 低 4 位是 BAR 类型标志d.bar0_sizeGetBARSize(bus,dev,func,0);// BAR2offset 0x18 in PCIe config spaceuint64_tbar2ReadPCIConfig64(bus,dev,func,0x18);d.bar2_basebar2~0xF;d.bar2_sizeGetBARSize(bus,dev,func,2);// 通过 BAR0 寄存器查询设备信息d.core_countReadDeviceRegister(d,REG_CORE_COUNT);d.hbm_sizeReadDeviceRegister(d,REG_HBM_SIZE);// NUMA 节点从 ACPI SLIT 表查或从 /sys/bus/pci 查d.numa_nodeGetNumaNode(bus,dev,func);devices.push_back(d);}}}returndevices;}关键设计PCIe 枚举不依赖操作系统驱动OS Driver 在此之前加载直接读 PCIe 配置空间MMIO 访问物理地址。因为 OS Driver 可能在驱动 binfmt 之后才加载而 CANN runtime 需要在 Docker 容器里直接用 NPU——绕开 OS Driver直接通过 VFIO/UIO 做用户态驱动。HBM 页表管理从虚拟地址到物理 HBM 页HBM 不是传统的 DDR 内存——它是高带宽存储器通过硅中介层interposer和 NPU Die 封装在一起。驱动需要管理 HBM 的物理页表分配和释放 2MB 大页。// driver/src/memory/hbm_allocator.cpp// HBM 页表管理 56GB HBM 的分配// 使用 buddy system伙伴系统 2MB 大页structHBMAllocator{staticconstexprintPAGE_SHIFT21;// 2MB 2^21 bytesstaticconstexprintPAGE_SIZE121;// 2,097,152 bytesstaticconstexprintMAX_ORDER16;// 2^16 × 2MB 128GB最大连续分配staticconstexprintMAX_NUM_PAGES;// HBM_SIZE / 2MB// 每阶链表order 0 2MB, order 1 4MB, ... order 16 128GBstd::vectorstd::listuint64_tfree_lists;// 每个 order 的空闲页链表std::vectoruint64_tpage_table;// 页表物理页编号 → order// 分配 size 字节的连续 HBMuint64_tAllocate(size_t size){// 1. 计算需要的最小 orderintorder0;size_t alloc_sizePAGE_SIZE;while(alloc_sizesize){alloc_size1;order;}// 2. 在 free_lists[order] 中查找intsearch_orderorder;while(search_orderMAX_ORDERfree_lists[search_order].empty()){search_order;}if(search_orderMAX_ORDER){// OOM没有足够大的连续块// 尝试碎片整理compact移动已分配的块释放大块Compact();search_orderorder;while(search_orderMAX_ORDERfree_lists[search_order].empty()){search_order;}if(search_orderMAX_ORDER){throwHBMOutOfMemory(Requested std::to_string(size) bytes);}}// 3. 从 higher order 分裂到 target orderuint64_tpagefree_lists[search_order].front();free_lists[search_order].pop_front();for(intosearch_order-1;oorder;o--){// 分裂order o1 的页分裂为两个 order o 的页uint64_tbuddypage(1ULL(oPAGE_SHIFT));free_lists[o].push_back(buddy);}page_table[pagePAGE_SHIFT]order;returnpage;}// 释放 HBM 页voidFree(uint64_tpage,intorder){page_table[pagePAGE_SHIFT]0;// Buddy merge如果 buddy 也是空闲的合并为更大的阶for(intoorder;oMAX_ORDER;o){uint64_tbuddypage^(1ULL(oPAGE_SHIFT));// 找 buddy 是否在 free_lists[o] 中autoitstd::find(free_lists[o].begin(),free_lists[o].end(),buddy);if(itfree_lists[o].end()){// buddy 不在空闲链表 → 不能合并 → 加入 free_lists[o]free_lists[o].push_back(page);return;}// 合并自己 buddy → order1free_lists[o].erase(it);pagemin(page,buddy);// 合并后的页起始地址}// 到达 MAX_ORDER放入最高阶链表free_lists[MAX_ORDER].push_back(page);}// 碎片整理compactvoidCompact(){// 移动已分配的块释放大的连续空间// 注compact 通常只在 OOM 时做代价高需暂停所有 Streamfor(intorderMAX_ORDER-1;order0;order--){for(autoitfree_lists[order].begin();it!free_lists[order].end();){uint64_tbuddy*it^(1ULL(orderPAGE_SHIFT));autobuddy_itstd::find(free_lists[order].begin(),free_lists[order].end(),buddy);if(buddy_it!free_lists[order].end()){uint64_tmergedmin(*it,buddy);free_lists[order].erase(it);free_lists[order].erase(buddy_it);free_lists[order1].push_back(merged);itfree_lists[order].begin();}else{it;}}}}}};为什么用 Buddy System 而不是 C mallocHBM 分配的是 GPU 内存需要物理连续且 2MB 对齐TLB 最小粒度。C malloc 的分配粒度是 8 bytes碎片小但地址不连续需要 IOMMU 映射NPU 不支持 IOMMU 重映射。Buddy System 天生对齐 2MB 边界且合并/分裂是 O(1) 的。DMA 命令提交从用户态到硬件队列// driver/src/dma/dma_engine.cpp// DMA 引擎提交 HBM↔Host 的数据搬运命令// 用户态CANN runtime通过 MMIO BAR0 寄存器直接提交命令structDMACommand{enumOpCode{H2D0x01,// Host → HBMCPU 内存 → NPU HBMD2H0x02,// HBM → HostNPU HBM → CPU 内存D2D0x03,// HBM → HBM同 NPU 内拷贝P2P0x04,// HBM → HBM跨 NPU通过 NVLink/HCCS};OpCode opcode;uint64_tsrc_addr;// 源地址Host 物理地址 或 HBM 偏移uint64_tdst_addr;// 目标地址uint32_tsize;// 搬运大小bytesuint32_tstream_id;// 关联的 Streamuint64_tevent_id;// 完成后触发的 Event};classDMAEngine{private:// 硬件队列环形缓冲区Ring Buffer// 基地址在 BAR0 寄存器中HOST_DMA_QUEUE_BASEstaticconstexprintQUEUE_DEPTH4096;// 4K 个命令staticconstexprintQUEUE_ENTRY_SIZE32;// 32 bytes per commandvolatileuint32_t*queue_base;// MMIO 映射的硬件队列uint32_tproducer_idx;// 软件维护的生产者指针volatileuint32_t*consumer_idx;// 硬件维护的消费者指针只读// 影子队列存放未提交的命令DMACommand shadow_queue[QUEUE_DEPTH];public:StatusInit(uint64_tbar0_base){// 步骤 1通过 MMIO 映射硬件队列// HOST_DMA_QUEUE_BASE 在 BAR0 的偏移 0x1000 处uint64_tqueue_physReadBAR0(bar0_base,0x1000);queue_base(volatileuint32_t*)MapPhysicalMemory(queue_phys,QUEUE_DEPTH*QUEUE_ENTRY_SIZE);// 步骤 2读取硬件维护的消费者指针寄存器uint64_tconsumer_regReadBAR0(bar0_base,0x1008);consumer_idx(volatileuint32_t*)MapPhysicalMemory(consumer_reg,4);producer_idx0;returnStatus::OK;}StatusSubmitDMA(constDMACommandcmd){// 步骤 1检查队列是否满uint32_tnext_producer(producer_idx1)%QUEUE_DEPTH;if(next_producer*consumer_idx){// 队列满 → 自旋等待硬件消费intretry0;while(next_producer*consumer_idxretry1000){__builtin_arm_isb();// 内存屏障ARM 架构retry;}if(retry1000){returnStatus::DMAQueueFull;}}// 步骤 2写入命令到影子队列shadow_queue[producer_idx]cmd;// 步骤 3写入硬件队列MMIO store// 命令格式32 bytes8 个 uint32// [0]: opcode | stream_id 8 | size 2// [1-2]: src_addr (64-bit)// [3-4]: dst_addr (64-bit)// [5]: event_id (32-bit)// [6-7]: reserveduint32_t*qentry(uint32_t*)(queue_baseproducer_idx*(QUEUE_ENTRY_SIZE/4));qentry[0]cmd.opcode|(cmd.stream_id8)|(cmd.size2);qentry[1]cmd.src_addr0xFFFFFFFF;qentry[2]cmd.src_addr32;qentry[3]cmd.dst_addr0xFFFFFFFF;qentry[4]cmd.dst_addr32;qentry[5]cmd.event_id;// 步骤 4推进生产者指针硬件看到后开始执行// 内存屏障确保命令写入完成后再推进指针__builtin_arm_dmb();// 数据内存屏障producer_idxnext_producer;// 步骤 5通知硬件doorbellWriteBAR0(bar0_base,0x1010,producer_idx);// HOST_DMA_DOORBELLreturnStatus::OK;}StatusWaitForEvent(uint64_tevent_id){// 轮询 Event 寄存器直到事件完成uint64_tevent_regbar0_base0x2000event_id*8;volatileuint64_t*event_ptr(volatileuint64_t*)MapPhysicalMemory(event_reg,8);while(*event_ptr!1){// 轮询等待生产环境用中断这里简化}// 清理事件*event_ptr0;returnStatus::OK;}};DMA 命令提交的关键路径MMIO store~150ns→ doorbell~30ns→ 硬件队列消费DMA 引擎读队列~200ns 启动延迟→ 实际传输1.2TB/s HBM 带宽。踩坑一PCIe BAR 大小不够导致的 MMIO 配置失败Ascend 910 的 BAR0 是 32KB——这是 PCIe 配置的 BAR 大小不能动态扩。如果用户态驱动注册了太多 MMIO 映射如 64 个 Stream × 32KB register space 2MB 32KBBAR0 映射失败 →torch.npu.init()抛 RuntimeError。// ❌ 每个 Stream 映射 32KB register space → BAR0 32KB 只能容纳 1 个 Streamfor(ints0;sMAX_STREAMS;s){uint64_tstream_basereg_bases*32768;// 32KB per streamRegisterMMIO(stream_base,32768);// 第 2 个 Stream 就失败}// ✅ 按需映射只有活跃的 Stream 才占用 BAR0 slot// 90% 的情况下活跃 Stream 4 → 4 × 32KB 128KB 32KB → 还是不够// 改进bar0 只是窗口实际映射用 BAR264GB HBM 映射// Stream 的寄存器空间映射到 HBM 地址BAR2bar0 只保留 1 个窗口解决方案stream register space 不映射到 BAR0——映射到 BAR2 的保留区域HBM 物理地址。bar0 只保留一个 4KB 窗口用于全局寄存器中断、设备状态。踩坑二Buddy System 的外部碎片和 OOM 假警报56GB HBM 分配成 1GB 2GB 1GB 52GB 后释放 2GB 那个块 → buddy 系统有 1GB 1GB 52GB连续但请求 3GB → 找不到连续 3GB 块1152 不连续需要 52GB 那个完整块分裂。OOM 假警报——实际上有 54GB 空闲。// ❌ 分配/释放模式导致外部碎片// 分配1GB (order 9), 2GB (order 10), 1GB (order 9)// 释放 2GB中间的 order 10 块// 空闲1GB (order 9) | 2GB (order 10, free) | 1GB (order 9) | 52GB (order 16)// ↑ 112GB但不在同一 buddy pair → 不能合并成 order 10// 申请 3GB → 找不到连续 order 11 → 失败// ✅ compact移动 1GB 块重组 buddy pairCompact();// 把 1GB 块移动到 52GB 区域的末尾// 空闲4GB (order 11) | 50GB (order 16)// 申请 3GB → 4GB 分裂 → 成功HBM compact 的代价暂停所有 DMA 传输H2D/D2H 暂停结束当前 kernel等最后一个 kernel 完成移动 HBM 块D2D 内部拷贝恢复 DMA。整个过程 ~5-10ms但只在 OOM 时触发。踩坑三轮询 Event 寄存器的 CPU 浪费上面的WaitForEvent用轮询spin-wait——如果 Event 在 100μs 后完成轮询了 100μs × 3GHz CPU 300K 个 cycle。如果有 32 个 Stream 都在等不同 Event → 32 × 100μs 的 CPU 时间 → CPU 占用 100%。// ❌ spin-waitCPU 空转 100μs浪费 CPU 时间while(*event_ptr!1){}// 3GHz CPU 100μs 300K 次读 300K 次分支预测// ✅ 中断驱动Edge-triggered MSI// 配置 MSI-X 中断向量voidSetupMSIX(intinterrupt_vector){// 写 MSI-X 表PCIe Capability 0x11uint64_tmsix_tableReadPCICapability(bus,dev,func,0x11);WriteMSIXEntry(msix_table,interrupt_vector,event_baseevent_id*8,// 中断产生时赋值为 10);// 中断清除时赋值为 0// 注册中断处理函数RegisterInterruptHandler(interrupt_vector,[event_id](){// 中断到达 → 设置 Event 1 → WaitForEvent 返回*event_ptr1;});}// ✅ 中断驱动的等待CPU 不空转StatusWaitForEventIRQ(uint64_tevent_id){if(*event_ptr1){*event_ptr0;returnStatus::OK;// 已经完成不用等}// 阻塞直到中断到达sem_wait(event_semaphores[event_id]);// 中断到达 → 自动唤醒*event_ptr0;returnStatus::OK;}中断驱动的等待CPU 不空转但这增加了上下文切换开销~5μs。100μs 的 DMA 等待 → 中断开销占比 5%。如果用轮询CPU 占用 100%。权衡短等待 10μs用轮询长等待 50μs用中断。driver 层从 PCIe 枚举直接读配置空间不依赖 OS 驱动→ HBM 页表Buddy System 2MB 大页合并/分裂 O(1)→ DMA 命令提交MMIO store 150ns → doorbell 30ns → 硬件队列启动 200ns。三个踩坑BAR0 32KB 映射不够用 → 映射到 BAR2 HBM 保留区、Buddy System 外部碎片 OOM 假警报 → compact 移动块重组 buddy pair5-10ms、spin-wait 轮询 Event 寄存器占满 CPU → 短等待轮询 长等待中断的混合策略。
