每组一个蛋白质$P$，原子数$I$，得到蛋白原子坐标矩阵$P=[p_1,p_2,\cdots ,p_I]$，形状$[I,3]$；

小分子$M$，原子数$J$，得到小分子原子坐标矩阵$M=[m_1,m_2,\cdots ,m_{y}J]$，形状$[J,3]$；

矩阵$P$和$M$做逐元素运算得到距离矩阵$D$，形状$[I,J]$：

$$\begin{gathered} D_{i,j}=\sqrt{(P_{i,x}-M_{i,x}{)}^2+(P_{i,y}-M_{i,y}{)}^2+(P_{i,z}-M_{i,z}{)}^2} \\ i\in [1,I],j\in [1,J] \end{gathered}$$

给定距离阈值$cutoff$，当距离小于$cutoff$时，选中该残基。

预设条件

硬件资源：

5060Ti，CUDA core 4608个，SM单元36个，显存16GB，带宽448GB/s

CPU 12700F，12核心，20线程，内存32GB

每个坐标值为float32类型，则每组$P$矩阵占用$3I\times 32=64I$；$M$矩阵占用$3J\times 32=64J$；$D$矩阵占用$32I\cdot J$。

原始任务

总任务数$T=1980$，大小$4.85$GB，设蛋白质平均原子数为5000个，小分子平均原子数为50个。

计算距离矩阵，每组需要进行3次减法、3次平方、2次加法、1次求根号

优化

文件多核心加载

CUDA批次打包

根号等价代换

距离是坐标差的平方和再开根号，可以等价为阈值的平方$cutof{f}^2$，用一次平方省去多次求根号操作，此时距离矩阵为：

$$D_{i,j}=(P_{i,x}-M_{i,x}{)}^2+(P_{i,y}-M_{i,y}{)}^2+(P_{i,z}-M_{i,z}{)}^2$$

每个核心需要进行3次减法、3次平方、2次加法

布尔比较

得到距离平方后，直接和阈值的平方$cutof{f}^2$比较，得到布尔值，并行得到布尔矩阵

对布尔矩阵按蛋白质的原子数聚合，得到大小为$J$的一维布尔向量

布尔向量再按残基聚合选中残基