1. 近内存计算系统的性能优化挑战
近内存计算(Compute-Near-Memory, CNM)架构正在重塑现代计算系统的设计范式。这种架构通过将处理单元直接嵌入或靠近内存芯片,从根本上改变了传统冯·诺依曼体系结构中处理器与内存分离的局限。在典型的CNM系统中,如UPMEM的DIMM设计,每个内存芯片集成了8个数据处理单元(DPU),每个DPU拥有专属的64MB DRAM存储区(MRAM)和64KB的SRAM缓存(WRAM)。这种紧密耦合的架构使得数据密集型应用如矩阵运算的访存延迟降低了约80%,同时能效比提升了3-5倍。
然而,这种架构创新也带来了独特的优化挑战。以UPMEM系统为例,其DPU采用14级流水线的RISC核心设计,支持24个硬件线程(tasklets)的并发执行。这种设计虽然提高了吞吐量,但也引入了复杂的资源竞争问题——当多个tasklets同时发起DMA请求时,由于共享的DMA引擎每次只能处理一个请求,会导致其他线程阻塞。我们的实测数据显示,在gemv(矩阵-向量乘)核函数中,不当的线程调度会使DMA冲突率高达35%,严重影响性能。
2. CoMoNM框架的架构设计
2.1 分层建模方法论
CoMoNM框架采用分层抽象的方法对CNM系统进行建模,其核心包含三个关键组件:
计算引擎模型:
- 精确模拟不同CNM架构的执行单元特性
- 支持流水线(如UPMEM的14级DPU流水线)和向量处理(如HBM-PIM的16个FPU阵列)两种模式
- 内置指令延迟表,包含200+种指令在不同架构下的时钟周期数
内存引擎模型:
class MemoryEngine { public: enum AccessType { BLOCKING, NON_BLOCKING }; void configure(uint32_t banks, uint32_t bank_size, AccessType type, uint32_t dma_channels); uint32_t estimateLatency(MemOp op); private: std::vector<MemoryBank> banks_; DmaController dma_; };该模型支持对多级内存层次(MRAM/WRAM/Register)的访问行为建模,特别是对非阻塞DMA传输的竞争场景进行精确模拟。
虚拟指令集(llvcnm):
- 抽象了UPMEM、HBM-PIM等不同CNM架构的指令集共性
- 包含LOAD/STORE/ALU等基础指令类型
- 通过后端适配器将通用指令映射到具体硬件实现
2.2 硬件无关的应用表示
CoMoNM的创新之处在于其硬件无关的中间表示设计。如图7所示的gemv核函数,可以通过多种前端语言描述:
MLIR表示:
linalg.generic { indexing_maps = [affine_map<(i,j)->(i,j)>, affine_map<(i,j)->(j)>, affine_map<(i,j)->(i)>], iterator_types = ["parallel", "reduction"] } ins(%A, %B : memref<8x16xi32>, memref<16xi32>) outs(%C : memref<8xi32>) { ^bb0(%a: i32, %b: i32, %c: i32): %prod = arith.muli %a, %b : i32 %sum = arith.addi %prod, %c : i32 linalg.yield %sum : i32 }UPMEM C代码:
#pragma tasklet void gemv(int* A, int* B, int* C, int M, int N) { for(int i=0; i<M; i++) { C[i] = 0; for(int j=0; j<N; j++) { C[i] += A[i*N+j] * B[j]; } } }
框架的cnm-gen组件会将这些高阶表示转换为统一的llvcnm虚拟汇编,如图10c所示。这个转换过程会保留关键的并行度信息和数据访问模式,为后续的映射优化奠定基础。
3. 迭代空间映射与优化
3.1 多维映射规范
CoMoNM引入的映射规范是其精准预估的核心。如图8所示的gemv映射示例,一个N维迭代空间在K级计算层次上的映射可表示为:
m = {[d₁¹,d₂¹,...,d_N¹], ..., [d₁^{K+1},d₂^{K+1},...,d_N^{K+1}]}其中每个元组表示在对应层次上各维度的划分因子。例如UPMEM的M3映射{(1,4),(2,4),(8,1),(12,32)}表示:
- 在第1级(Rank层)沿第2维划分4份
- 在第2级(DPU层)沿第1维划分2份、第2维4份
- 最终每个tasklet处理12×32的子空间
3.2 非直观优化发现
通过系统的设计空间探索,CoMoNM能够发现违反直觉但高效的映射方案。在gemv测试中,我们发现:
UPMEM最佳配置:
- 传统认知:增大tasklet数量(接近24)可提高并行度
- 实际发现:12-16个tasklet时DMA冲突最少
- 原因:适度的并行度允许更好的DMA请求交错
HBM-PIM指令排序:
// 低效版本(延迟:128周期) READ E,R0,C ADD E,R1,C WRITE E,R2,C READ O,R0,C ADD O,R1,C WRITE O,R2,C // 优化版本(延迟:64周期) READ E,R0,C READ O,R0,C ADD E,R1,C ADD O,R1,C WRITE E,R2,C WRITE O,R2,C交替访问偶/奇存储体的指令序列可使吞吐量提升2倍,这是通过CoMoNM的自动探索发现的。
4. 性能预估引擎实现
4.1 分层抽样算法
如算法1所示,CoMoNM采用创新的分层抽样方法进行快速预估:
- 从完整迭代空间中提取代表性子集(get_portion)
- 在模型上精确模拟这些子集的执行(simulate)
- 根据子集间的性能变化趋势外推整体性能
这种方法的优势在于:
- 对大规模问题(如2048×2048矩阵)只需模拟512×512子集
- 保持<3%的预估误差的同时,速度比全模拟快10^7倍
4.2 竞争建模
内存访问竞争是影响预估精度的关键因素。CoMoNM的mem-engine实现了精确的竞争检测:
def accept(self, instruction): if instruction.type == DMA_REQUEST: if self.dma.busy: self.stall_cycles += estimate_contention_delay() return False self.dma.start(instruction) return True该模型会跟踪DMA引擎状态,当检测到多个tasklet同时发起请求时,自动插入适当的停顿周期。
5. 集成与评估
5.1 编译器集成
CoMoNM可与现有CNM编译流无缝集成。图1展示了其在CINM编译器中的集成点:
- 高层优化阶段:指导循环分块策略
- 映射决策阶段:评估不同卸载方案
- 代码生成阶段:验证指令调度效果
实测显示,集成CoMoNM后编译器做出的自动决策,比人工优化方案平均性能提升22%。
5.2 精度验证
我们在UPMEM真实硬件和三星HBM-PIM模拟器上验证了框架精度:
| 基准测试 | 实测值(ms) | CoMoNM预估(ms) | 误差率 |
|---|---|---|---|
| gemv-1024 | 14.2 | 14.8 | +4.2% |
| spmv-4096 | 68.5 | 65.9 | -3.8% |
| conv2d-128 | 112.3 | 115.6 | +2.9% |
关键发现:
- 对于规则访问模式(如gemv),误差<5%
- 不规则模式(如spmv)因竞争波动稍大,但仍<8%
- 比传统模拟器快7个数量级,使实时优化成为可能
6. 实际应用建议
基于我们的实施经验,给出以下CNM优化建议:
DMA最佳实践:
- 将大传输拆分为多个64KB块(UPMEM WRAM大小)
- 在计算密集区预取下一块数据
- 避免在临界区集中发起DMA请求
线程配置经验:
// 次优:使用全部24个tasklet #pragma tasklet num=24 void kernel() { ... } // 优化:根据负载选择12-16个tasklet #pragma tasklet num=14 void optimized_kernel() { ... }数据布局优化:
- 对HBM-PIM采用交错存储(图3e)
- UPMEM中确保MRAM数据对齐64字节边界
这套方法论已在多个AI加速项目中验证,相比传统GPU方案,在推荐系统embedding查询等场景可实现3.8倍能效提升。未来我们计划将CoMoNM扩展到更多新兴CNM架构,如Alibaba HB-PNM和SK hynix AiM,持续推动近内存计算生态发展。