GPU加速的REPEAT电荷计算程序

slxc920113

最近在阅读gpu4pyscf的代码，学习GPU编程，就拿REPEAT电荷这个比较耗时的程序练手了。核心的REPEAT计算算法并不复杂，就是分成实空间+倒空间，进行Ewald求和，熟悉MD程序的应该都懂。起初用Python重写的程序只有300多行，但是纯CPU下计算很慢，至少比C++版本慢了两个数量级。然后就把所有的Python加速手段一点点往上面叠加，最后代码量到了离谱的2000多行。

1. GIL版本的Python， 多进程
2. GIL版本的Python， 多进程+numba jit计算加速
3. GIL版本的Python，numba jit +TBB 并行
4. noGIL版本的Python（3.13+），多线程
5. 采用Cupy或者Pytorch进行GPU加速
6. 采用Cupy或者Pytorch进行GPU加速+多线程调度多张显卡
7. FP64精度降低到FP32混合精度

最终测试下来在7950X+4090上，一个200原子的MOF体系的REPEAT电荷计算python版本可以压缩到10秒以内，纯C++的CUDA并行版本大约3-4秒。
欢迎试用：LZ0211/pyREPEAT: Python program for calculate Electrostatic Potential fitted charges for periodic systems

1. python repeat.py [CUBE_FILE] [OPTIONS]
   
2. 
   
3. 位置参数:
   
4.   cube_file             输入的CUBE文件路径
   
5. 
   
6. 可选参数:
   
7.   --fit-type {0,1}      拟合类型: 0=不考虑周期性, 1=考虑周期性 (默认: 1)
   
8.   --vdw-factor FLOAT    VDW半径缩放因子 (默认: 1.0)
   
9.   --vdw-max FLOAT       VDW半径上限 (默认: 1000.0)
   
10.   --cutoff FLOAT        实空间截断半径 (默认: 20.0)
    
11.   --total-charge FLOAT  体系总电荷 (默认: 0.0)
    
12.   --cores N             CPU核心数 (默认: 自动检测)
    
13.   --gpu [DEVICES]       启用GPU加速
    
14.                         示例: --gpu, --gpu 0, --gpu 0,1,2, --gpu all
    
15.   --stats               计算统计量(RMS误差)。默认关闭以节省时间
    
16.   --fp64                强制双精度fp64 (默认fp32)
    
17.   --block-k N           k空间分块大小 (默认: 64)
    
18.   --symm-file PATH      对称性约束文件
    
19.   --resp-file PATH      RESP参数文件
    
20.   --qeq-file PATH       QEq参数文件
    
21.   --charge {repeat,qeq} 电荷模型 (默认: repeat)
    
22.   --output PATH         输出文件路径

复制代码

repeat.py (107.65 KB, 下载次数 Times of downloads: 5)

2026-2-3 16:16 上传 Uploaded

点击下载
Click to download