mat-chem-sim-pred 把昇腾 NPU 用在了传统高性能计算(HPC)的典型场景——分子动力学模拟、量子化学计算、材料性能预测。这几个领域的共性是用数学方程描述原子和分子之间的相互作用,大量矩阵运算天然适合 NPU 加速。
三类仿真任务
| 任务 | 数学本质 | 矩阵规模 | 典型耗时(CPU 64核) |
|---|---|---|---|
| 分子动力学(MD) | 牛顿运动方程积分(N² 粒子对计算) | 100K-1M atoms | 数小时-数天 |
| DFT 电子结构 | 自洽场迭代(对角化 Fock 矩阵) | 1000-5000 basis | 数分钟-数小时 |
| 材料性能预测 | 晶体图神经网络推理 | 100-1000 atoms/cell | 数秒 |
MD 是三者中算量最大的——每个时步都要算每个原子对之间的作用力(O(N²) 粒子对),然后做时域积分。NPU 的 Cube 单元在这里直接发挥作用:粒子对作用力计算是矩阵乘的天然应用。
分子动力学的 Lennard-Jones 势函数
两个原子之间的势能与距离的关系由 Lennard-Jones 势函数描述:
V(r) = 4ε × [(σ/r)^12 - (σ/r)^6] r 是原子间距离 ε 是势阱深度(化学元素相关) σ 是碰撞直径(化学元素相关) r^12 项:排斥力(原子靠太近时推开) r^6 项:范德华吸引力计算每个原子的受力需要先算势能的梯度:F_i = -∂V/∂r_i。对所有原子对求和。
// mat-chem-sim-pred/kernels/lennard_jones.cpp__aicore__voidLennardJonesForce(GlobalTensor<float>&forces,// 每个原子的受力 [N, 3] (x,y,z)GlobalTensor<float>&positions,// 每个原子的位置 [N, 3]GlobalTensor<float>&epsilon,// 每种元素的势阱深度 [element_types]GlobalTensor<float>&sigma,// 每种元素的碰撞直径 [element_types]GlobalTensor<int>&atom_type,// 每个原子属于哪种元素 [N]intN,// 原子总数floatcutoff// 截断半径(超出付distance 视为 0)){// 初始化力为零forces=0.0f;// 第一步:计算所有原子对的距离矩阵 D_ij// D_ij = ||pos_i - pos_j||// N 个原子的距离矩阵是 N×N → 1M atoms 时距离矩阵是 1M×1M// 直接计算需要 10^12 次距离计算(不可行)// 用邻居列表(neighbor list)削减计算量:// 只计算距离 < cutoff 的原子对对// cutoff = 10Å 时每个原子的邻居数 ≈ 50-200(取决于密度)for(inti=0;i<N;i++){floatsigma_i=sigma[atom_type[i]];floateps_i=epsilon[atom_type[i]];floatfx=0.0f,fy=0.0f,fz=0.0f;floatxi=positions[i*3];floatyi=positions[i*3+1];floatzi=positions[i*3+2];// 只遍历邻居列表中的原子对(不含自身)for(intj:neighbor_list[i]){floatxj=positions[j*3];floatyj=positions[j*3+1];floatzj=positions[j*3+2];// 距离向量floatdx=xi-xj;floatdy=yi-yj;floatdz=zi-zj;// 距离平方(避免开方)floatr2=dx*dx+dy*dy+dz*dz;// 截断判定floatsigma_ij=0.5f*(sigma_i+sigma[atom_type[j]]);floateps_ij=sqrt(eps_i*epsilon[atom_type[j]]);floatcutoff2=(cutoff*sigma_ij)*(cutoff*sigma_ij);if(r2>cutoff2)continue;// Lennard-Jones 力函数// 把距离缩放:r2_scaled = (σ_ij / r)^2floatsigma2_over_r2=(sigma_ij*sigma_ij)/r2;// (σ/r)^6 和 (σ/r)^12floats6=sigma2_over_r2*sigma2_over_r2*sigma2_over_r2;floats12=s6*s6;// 力的大小 = 24ε_ij / r * (2*(σ/r)^12 - (σ/r)^6)floatinv_r=sqrt(1.0f/r2);floatforce_mag=24.0f*eps_ij*inv_r*(2.0f*s12-s6);// 力分解到三个方向fx+=force_mag*dx*inv_r;fy+=force_mag*dy*inv_r;fz+=force_mag*dz*inv_r;}forces[i*3]=fx;forces[i*3+1]=fy;forces[i*3+2]=fz;}}邻居列表每 10-20 个时步重建一次(原子移动后邻居关系变化)。NPU 上用 SDMA 把邻居列表的数据搬到 L1——不需要从 HBM 逐个加载原子位置。
邻居列表构建
邻居列表的构建是 MD 仿真的关键数据结构——决定了计算量的 O(N²)→O(N×neighbors) 削减。
// mat-chem-sim-pred/kernels/neighbor_list.cpp// Verlet 邻居列表法:// 把空间分成三维网格(每个格点大小 = cutoff + skin)// 只有相邻 27 个格点内的原子才可能相互作用__aicore__voidBuildNeighborList(GlobalTensor<int>&neighbor_list,// [N, max_neighbors]GlobalTensor<int>&neighbor_count,// [N]GlobalTensor<float>&positions,// [N, 3]floatcell_size,// 格点大小intgrid_x,intgrid_y,intgrid_z,// 三维网格尺寸intN){// 第一步:分配原子到格点// grid_id = (x/cell_size) + (y/cell_size)*grid_x + (z/cell_size)*grid_x*grid_yLocalTensor<int>grid_atoms[grid_x*grid_y*grid_z];// 每个格点的原子列表for(inti=0;i<N;i++){intgx=int(positions[i*3]/cell_size);intgy=int(positions[i*3+1]/cell_size);intgz=int(positions[i*3+2]/cell_size);// 边界检查gx=Clamp(gx,0,grid_x-1);gy=Clamp(gy,0,grid_y-1);gz=Clamp(gz,0,grid_z-1);intgid=gx+gy*grid_x+gz*grid_x*grid_y;Append(grid_atoms[gid],i);}// 第二步:遍历每个格点及其 27 个相邻格点for(intgz=0;gz<grid_z;gz++){for(intgy=0;gy<grid_y;gy++){for(intgx=0;gx<grid_x;gx++){intgid=gx+gy*grid_x+gz*grid_x*grid_y;int*atoms_in_cell=grid_atoms[gid];intcount=SizeOf(atoms_in_cell);// 遍历 27 个相邻格点(包括自己)for(intdz=-1;dz<=1;dz++){for(intdy=-1;dy<=1;dy++){for(intdx=-1;dx<=1;dx++){intnx=gx+dx,ny=gy+dy,nz=gz+dz;if(IsOutOfBounds(nx,ny,nz))continue;intnid=nx+ny*grid_x+nz*grid_x*grid_y;int*neighbor_atoms=grid_atoms[nid];// 当前格点 vs 相邻格点:全部原子对都加入邻居列表for(inta:atoms_in_cell){for(intb:neighbor_atoms){if(a<=b)continue;// 只计算一次(a > b)floatdx=positions[a*3]-positions[b*3];floatdy=positions[a*3+1]-positifloatdz=positions[a*3+2]-positions[b*3+2];floatr2=dx*dx+dy*dy+dz*dz;if(r2<cutoff2){neighbor_list[a*max_neighbors+neighbor_count[a]]=b;neighbor_count[a]++;}}}}}}}}}}邻居列表构建本身是 O(N) 的(每个原子只检查 27 个相邻格点),比 O(N²) 的直接搜索快得多。格点大小选得好(cell_size = cutoff + skin),可以保证不遗漏任何邻居对。
自洽场迭代(DFT 电子结构计算)
密度泛函理论(DFT)的电子结构计算是量子化学的基础。核心是自洽场(SCF)迭代——反复构造和対角化 Fock 矩阵:
初始猜一个电子密度 ρ(r) ↓ 构造 Fock 矩阵 F[ρ](Kohn-Sham 方程) ↓ 対角化:F C = ε S C(广义特征值问题) ↓ 更新电子密度 ρ(r) = Σ |φ_i(r)|² ↓ 检查收敛:|ρ_new - ρ_old| < tolerance矩阵规模受 basis set 大小限制(1000-5000),问题是自洽场迭代可能需要 20-50 次——每次都要对角化一次矩阵。対角化本身是 O(N³),NPU 的 Cube 单元在这里大幅加速。
// mat-chem-sim-pred/kernels/scf_iteration.cpp__aicore__boolSCFIteration(GlobalTensor<float>&Fock,// Fock 矩阵 [B, B]GlobalTensor<float>&S,// 重叠矩阵 [B, B](固定不变)GlobalTensor<float>&C,// 分子轨道系数 [B, B](输出)GlobalTensor<float>&eigenvalues,// 轨道能量 [B]GlobalTensor<float>&density,// 密度矩阵 [B, B]intB,// basis set 大小floattolerance){boolconverged=false;LocalTensor<float>density_old(B*B);density_old=0.0f;for(intiter=0;iter<50;iter++){// 第一步:广义特征值分解 F C = ε S C// NPU 上用分治法做广义对称特征值分解GeneralizedEigenSolver(Fock,S,C,eigenvalues,B);// 第二步:用前 50% 的占据轨道构造新密度矩阵intn_occ=B/2;// 占据轨道数LocalTensor<float>C_occ(B,n_occ);SliceColumns(C,C_occ,0,n_occ);// 密度矩阵 D = C_occ × C_occ^T// [B, n_occ] × [n_occ, B] → [B, B]// Cube 单元:稠密矩阵乘MatMul(density,C_occ,Transpose(C_occ));// 第三步:收敛检查LocalTensor<float>diff(Axpy(density,-1.0f,density_old));floatrmsd=RMSNorm(diff);if(rmsd<tolerance){converged=true;break;}// 第四步:用新密度构造 Fock 矩阵// F[ρ] = H_core + J[ρ] - 0.5 * K[ρ]// J 是库仑积分(经典电子排斥)// K 是交换积分(量子力学交换项)BuildFockMatrix(Fock,density,H_core,ERIs,B);// 阻尼:Fock = α·F_new + (1-α)·F_old// 防止迭代震荡floatalpha=0.7f;Scale(density_old,(1-alpha));Axpy(density_old,alpha,density);}returnconverged;}DFT 计算的对角化和矩阵构造在 5000 basis 规模下,NPU 比 64 核 CPU 快 5-8 倍。关键是 SCF 迭代的次数——NPU 上每步 50ms vs CPU 上每步 300ms,20 次迭代就是 1s vs 6s 的差距。
踩坑一:截断半径邻居列表遗漏原子对
邻居列表的格点大小 cell_size = cutoff + skin。skin 是安全缓冲——如果原子移动速度太快,可能从一个格点跳到相邻格点之外,导致原本在邻居列表外的原子对「靠上来」了。Verlet 邻居列表算法用 skin 解决:每次重建邻居列表时把列表扩大到 cutoff + skin,后面几个时步即使原子移动了一些,原本的邻居仍然在列表内。
但 skin 不能太大——太大会包含太多无用的原子对,把 O(N×neighbors) 膨胀回接近 O(N²)。
正确设置:
// 经验值:skin ≈ 0.3 × cutofffloatcutoff=10.0f;// Åfloatskin=2.0f;// Å(20% cutoff)floatcell_size=cutoff+skin;// 12.0 Å// 每 15 个时步重建邻居列表// 原子在 15 步 × 0.5fs/步 = 7.5fs 内的位移 ≈ 0.1-0.5 Å(小于 skin)// 所以不会遗漏邻居if(step_count%15==0){BuildNeighborList(...);}踩坑二:DFT 的密度矩阵数值对称性丢失
密度矩阵 D 在理论上是对称矩阵(D_ij = D_ji)。经过多次矩阵乘和迭代更新的数值累积后,矩阵的非对称部分(由浮点舍入引入)逐渐积累,导致対角化时特征值出现虚部。
修复:每次构造密度矩阵后强制对称化。
// 密度矩阵对称化:D = 0.5 × (D + D^T)MatMul(density,C_occ,Transpose(C_occ));// 强制对称:不是简单的 D = (D + D^T) / 2// NPU 上用一把 Vector 指令完成LocalTensor<float>d_transpose(B,B);Transpose(density,d_transpose);// 注意:不是直接 Copy,是转置Add(density,density,d_transpose);Scale(density,0.5f);踩坑三:MD 模拟的时步选择
时步太小(0.1fs)计算量大,太大(10fs)原子运动急剧变化式中数值发散。化学键的振动周期是 ~10fs,合理时步是 0.5-1.0fs。
Shake 约束算法:对含氢原子的键(包括 O-H、N-H、C-H 等)施加固定键长约束,允许时步增大到 2fs。
// Shake 算法:动力学中保持化学键键长不变// 氢原子的质量小(1 amu)、振动频率高(~3000 cm⁻¹)// 如果把 H-X 键锁定,时步可以从 0.5fs 增到 2.0fs(4×)voidShakeConstraints(GlobalTensor<float>&positions,// 所有原子位置(更新后)intN_bonds,// 约束的键的数量floattarget_bond_length// 键的目标长度){// 对每个约束的键(X-H bond)for(intb=0;b<N_bonds;b++){inta1=bonds[b].atom1;// 通常是 Hinta2=bonds[b].atom2;// 通常是 O/N/C// 当前键长floatdx=positions[a2*3]-positions[a1*3];floatdy=positions[a2*3+1]-positions[a1*3+1];floatdz=positions[a2*3+2]-positions[a1*3+2];floatcurrent_length=sqrt(dx*dx+dy*dy+dz*dz);// 拉格朗日乘子法:调整原子位置使键长恢复到目标值// λ = (d² - r²) / (2 × r × (1/m1 + 1/m2) × Δt²)floatd=target_bond_length;floatinv_mass_sum=1.0f/mass[a1]+1.0f/mass[a2];floatlambda=(d*d-current_length*current_length)/(2.0f*current_length*inv_mass_sum*dt*dt);// 移动原子(质量加权)positions[a1*3]-=lambda*dx*(1.0f/mass[a1]);positions[a1*3+1]-=lambda*dy*(1.0f/mass[a1]);positions[a1*3+2]-=lambda*dz*(1.0f/mass[a1]);positions[a2*3]+=lambda*dx*(1.0f/mass[a2]);positions[a2*3+1]+=lambda*dy*(1.0f/mass[a2]);positions[a2*3+2]+=lambda*dz*(1.0f/mass[a2]);}}Shake 约束把时步从 0.5fs 增加到 2fs——模拟 1ns 需要 5 亿步变成 1.25 亿步——总计算时间缩短到原来的 1/4。
mat-chem-sim-pred 把 NPU 用在了与传统 AI 完全不同的场景:分子动力学、量子化学、材料模拟。算子的本质是矩形的矩阵乘(Lennard-Jones 力函数)、特征值分解(DFT 自洽场)、邻居搜索(Verlet 列表)——都是传统高性能计算的经典算法。NPU 的价值超越深度学习——在科学计算的 HPC 领域同样成立。