Uni-Dock论文详解：以C++/CUDA优化的视角

引子

赶着末班车报名了Deepbrew，得知要写个Notebook，搜了一下发现关于Uni-Dock大家已经写得差不多了，甚感欣慰的同时也不免有种“前人之述备矣”的感觉——Uni-Dock要跑起来，的确就应该是两三个Notebook就要交代清楚的事，否则就是没有做到足够易用。但是，那我还能写什么呢？更广泛地说，能在毕业出国前留下点什么呢？

想来想去，就写写开发Uni-Dock过程中的故事吧！这个项目从21年一路走来，起初只是一个“用GPU打败CPU上的Vina”的蓝图，又经过22年初一人在北大的宿舍debug大半个寒假，接着到终于能跑通100倍的加速，再到后来文章成稿、投稿、拒稿、改了再投、再修，Uni-Dock的性能越来越好，功能越来越多，参与的小伙伴们也越来越多，甚至现在有了一定的影响力。这个项目也贯穿了我高年级本科阶段，在每一门课上我应该都写过Uni-Dock的代码，申请PhD的时候SoP也是基本围绕它展开的，似乎我也在被这个项目定义着，以至于在某些不起眼的瞬间决定了我的未来。

Uni-Dock的核心，说来太过简单，一开始就只是加速而已。甚至都没有底层算法的创新，在一通计算后发现用GPU的一个thread如果代替CPU的一个核心，估摸着能加速最高上万倍，然后就开干了。过程中才发现，完全没有这么简单，首先要能赶上Vina-GPU，把CUDA的Monte-Carlo/MFGS搜索跑在GPU上，然后再用批处理提高并行度，最后看着Nsight Compute一个一个用tricks解决热点，榨干GPU。这三步也就是论文中的三个阶段（Stage），会在本Notebook中做详细介绍。之后基本上能在速度上达到要求，则需要详尽、扎实的各种测试，说服别人这是可靠的，并且一步步工程化，提高易用性。这一部分的图表基本上在论文中，也会在本Notebook中提到。最后，在Hermite平台上线，持续开发各个feature，进入稳定的维护和投稿阶段。其中Biased Docking、SDF support、Uni-Mol + Uni-Dock、网页APP前端就是这段时间做出来的，也有Notebook已经写到了。

最近，有新的小伙伴加入分子对接的开发，关于watvina、新的Constrain Docking、新的打分函数，有更多的人在期待，有更贴近应用的需求在出现。外部的合作交流接踵而至，Github issue接连不断，这个曾经单兵作战的项目迫切需要越来越多开发者的合力去维护，去做更多事，发挥更大的价值。如何将开发过程中的细节交代清楚，让未来的同学接下这个接力棒也成为我最近思考的事——而这，必然涉及到对Uni-Dock源代码的深入理解。换言之，就是要以C++/CUDA优化的视角去解析Uni-Dock的开发核心，将三个阶段的优化以更浅显的方式对着代码一一解释，将后来的各种Memory优化和新feature的方法论和硬性约束理清——最后把能做的都做好，不能做的静静等待某个灵光一现的基础创新。

那么，就开始吧（未完待续）

Baseline: Autodock Vina & Vina-GPU

这两个项目作为Uni-Dock开发前就已经release。Autodock Vina在业内广泛使用，有非常多的工作以此为基准，也招致了许多批评：打分函数筛选性能不行、代码结构宛如*山、层出不穷的内部报错……但是，Vina（版本1.2)确实是当下可用的开源框架里最具有平台属性的，它继承了Autodock4的pdbqt格式、打分函数和map+搜索的对接模式。下面首先对Vina的代码进行力所能及的梳理（建议对照源代码阅读）：

对于设定参数不过多赘述。在main.cpp中，首先根据参数完成搜索模式（是否local、是否score_only）、蛋白配体对象、打分函数和map的构建。

对于对接的蛋白配体和一些读写接口，以及主要是搜索之外的相关参数，Autodock Vina内部用Vina类进行梳理。对于存储蛋白配体复合物，用model类统一管理，其中又分为蛋白信息和配体信息；蛋白读入后刚性部分实际上不需要存储键的信息，所以直接当作离散的原子即可（model.grid_atoms)，生成的格点图Gridp Map存储在cache或ad4cache类中。配体需要存储原子信息和键的信息，以及中心原子ROOT，在搜索的不同阶段分别位于ligand,conf,ligand_conf类中，包含暂时的构象信息，最后是蛋白柔性残基，约等于固定了一端的配体，实际上残基和配体都是类似于树的结构，如图所示：

对于Monte-Carlo搜索，用mc类及并行的parallel_mc类统一管理，而对于其中的BFGS拟牛顿算法，有若干类和函数模板相互嵌套（quasi_newton,line_search()），因此代码比较难懂。不过，最关键的是global_search函数和monte_carlo::operator()，看懂后基本能理解搜索流程。最后是整理结果并输出，需要注意的是微调的过程（refine）是直接根据蛋白原子进行能量计算，跳过了格点图。

Vina并行加速的关键在于parallel_mc类中利用boost库进行了CPU层面的多线程并行，且在线程内部用了比较高效的MC/BFGS算法。而Vina-GPU则是基于Vina做了GPU的适配，仅仅将MC/BFGS算法迁移至OpenCL框架下的GPU就能达到数倍的加速，在此就简单展示：

Uni-Dock三个阶段的优化

第一阶段：单配体多构象并行搜索

在优化之前，我们可以发现热点集中在eval_deriv()函数，它是用于在搜索中计算当前构象的能量和梯度的函数 它的调用者就是之前所说的monte_carlo::operator()等函数： 我们将这部分时间统一算作全局搜索（global search）的时间，占据了整个Wall time的95%以上。这一阶段的优化集中在将global search全部迁移至GPU上，利用GPU的并行能力进行加速。 然而，第一次优化后使用和CPU相同的搜索步数和宽度，却发现需要执行40分钟以上。只有个将搜索步数大幅度减小，并加大宽度，才能比CPU有加速。究其原因，是因为GPU每个thread的时间比CPU长。

第二阶段：多配体并行分子对接

第一阶段中的加速仍然和CPU没有拉开差距。GPU利用率也非常低下，为了进一步并行化，我们使用多配体来填满GPU的可用线程。

第三阶段：“榨干”GPU计算性能