用 MPI 实现高性能分布式程序

COMP90025 是墨大的并行与多核计算（Parallel and Multicore Computing）。这门课有两个大作业，合起来正好是一个简单的并行拼写检查器。第一个作业用的是 OpenMP，从文档里收集不重复的单词，然后第二个作业就换成 MPI，把这些单词拿去和词典对一遍，看看哪些拼错了，又该改成什么。

这篇是第二个作业做完以后写的，所以文里的代码和数据都是最终提交的那一版，代码放在 xxxbrian/mpi-spell-checker。

题目

程序的用法很简单。

./spellcheck dictionary.txt word_list.txt

dictionary.txt 是排好序的词典，word_list.txt 是要检查的单词。对词表里的每个词，如果它不在词典里，就要生成它所有可能的候选，也就是改一个字符、删一个字符或者加一个字符，这样得到的新词里，凡是在词典里的，都是它可能的正确拼写。另外还有一条关于大小写的规则，一个词的首字母换成小写之后如果在词典里，那它也算在词典里。

输出里每个拼错的词占一行，后面跟着它的候选，像这样。

appel:  appeal apple rappel
stap:  sap snap stab step stop strap tap

行的顺序是按候选的数量从少到多排的，数量一样的话，就按字典序。

题目还加了一个限制：每个进程只能保留词典的一部分，不能各自存一整份。课程给了一份单线程的参考实现，用的是 BK-tree，而我们的程序则要放到学校的 HPC 集群 Spartan 上去跑。成绩分成代码和报告两块，因为报告占的分数更高，所以实验和分析也不能马虎。

先换一个算法

动手之前，我先把参考实现跑了一遍。它用 BK-tree 在词典里找编辑距离不超过 1 的词，在最大的一组数据上，词典有 62 万多个词，词表有 65 万多个词，结果它要跑大约 32 分钟。

我的做法就简单得多，既然编辑距离固定是 1，那就直接把所有候选词都生成出来，再去哈希表里一个个查。生成候选的代码并不复杂。

std::vector<std::string> GenerateEditDistance1Candidates(const std::string& word) {
  std::vector<std::string> out;
  const std::string alphabet = "abcdefghijklmnopqrstuvwxyz"
                               "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
                               "0123456789";

  for (size_t i = 0; i < word.size(); ++i)  // 删除
    out.push_back(word.substr(0, i) + word.substr(i + 1));

  for (size_t i = 0; i < word.size(); ++i)  // 替换
    for (char c : alphabet)
      if (c != word[i]) {
        std::string s = word;
        s[i] = c;
        out.push_back(s);
      }

  for (size_t i = 0; i <= word.size(); ++i)  // 插入
    for (char c : alphabet)
      out.push_back(word.substr(0, i) + c + word.substr(i));

  return out;
}

一个长度为 L 的词，会生成大约 124 × L 个候选，所以一个八九个字母的词，候选也就是一千上下。这个数量对哈希表来说完全不算什么，于是单个进程跑下来只要 20 秒左右，其中 19 秒花在生成候选和查询上。

在这组数据里，一共有 30627 个词是拼错的，每个词一千多个候选，加起来要查三千多万次。

所以说，快了差不多一百倍的这一步其实和并行没有关系，只是换了一个算法而已。那么剩下的问题就是，怎么把这 20 秒分给多个进程。

最早的版本

最早的版本我写得很直接，目标就是让它先跑起来。

词典是每个进程各自读进来，然后挑出属于自己的那一部分，按哈希分片。要判断一个词在不在词典里，就对每个词做一次 MPI_Allreduce，让所有进程都查一遍自己的分片，再把结果做逻辑或。

for (size_t i = 0; i < word_list.size(); ++i) {
  bool local_hit = local_dict.Contains(word_list[i]);
  bool global_hit = false;
  MPI_Allreduce(&local_hit, &global_hit, 1, MPI_C_BOOL, MPI_LOR, MPI_COMM_WORLD);
  in_dict[i] = global_hit;
}

候选词的验证也是类似的思路，把候选发给每一个进程，让它们各自查，再把命中的发回来，发送和接收都是阻塞的。

这么写思路上没有问题，但是毛病也很明显。词表有 65 万个词，就意味着要做 65 万次集合通信。每个候选都要发给所有进程，数据量也跟着变大。更麻烦的是阻塞的发送和接收会死锁，两个进程同时给对方发一大块数据，谁都发不出去，程序就卡在那里了。

死锁是最先处理的。一开始我只是把交换改成了 MPI_Sendrecv，后来干脆把整个验证流程改成了请求和响应。等回复的时候不能干等着，因为对方可能也在等我，所以每次循环都用 MPI_Iprobe 看一眼有没有别人发来的请求，有的话就顺手处理掉。

std::vector<int> WaitForResponse(int target, const DistributedDict& dict) {
  while (true) {
    DrainIncomingRequests(dict);  // 等回复的时候，也在帮别人查
    int flag = 0;
    MPI_Iprobe(target, kResponseHeader, MPI_COMM_WORLD, &flag, MPI_STATUS_IGNORE);
    if (flag)
      return ReceiveResponse(target);
  }
}

这样一来，等待回复时仍然会处理其他进程的请求，避免大家只等回复、却不回答请求的僵局。

然后是读文件的问题。既然题目要求词典是分布式的，我就改成了用 MPI-IO，每个进程只读自己那一段。不过文件是按字节切的，切开的地方很可能落在一个词的中间，于是起点和终点都要往后挪到下一个换行，这样每个词就只会属于一个进程。

MPI_Offset chunk = file_size / world_size;
MPI_Offset start = rank * chunk;
MPI_Offset end = (rank == world_size - 1) ? file_size : (rank + 1) * chunk;

if (rank > 0)              start = SkipToNextLine(fh, start);
if (rank < world_size - 1) end = SkipToNextLine(fh, end);

MPI_File_read_at(fh, start, buffer, end - start, MPI_CHAR, MPI_STATUS_IGNORE);

上面是示意代码，省略了缓冲区分配、EOF 和大文件计数处理。实际找行边界时，是用 MPI_File_read 一个字节一个字节往后读，直到遇到换行为止。

词典怎么分

最后的版本里，每个词都有一个主人，也就是这个词的哈希值对进程数取模得到的那个进程。词典里的词会被发给各自的主人，主人把它们放进自己的哈希表里。

这里有个小细节，因为首字母大小写的规则，Apple 和 apple 必须由同一个进程来处理。所以我在算哈希之前，会先把首字母转成小写，这样不管是哪种写法，都会落到同一个进程。

int GetOwnerRank(const std::string& word, int world_size) {
  std::string key = word;
  if (!key.empty())
    key[0] = std::tolower(static_cast<unsigned char>(key[0]));  // 首字母统一转小写
  return std::hash<std::string>{}(key) % world_size;
}

上面的代码只展示归属规则，报告中的实现使用 FNV-1a 计算哈希。

同样的规则对词表也适用。一个词的主人，既负责判断它在不在词典里，也负责在它拼错的时候生成候选。词表在每个进程里都有一份，每个进程按下标轮流负责其中一部分，把词发给它的主人去查。因为查的人和之后要用结果的人是同一个，所以结果也就不用再广播了。顺带一提，原来那 65 万次 Allreduce 在这里变成了一次 Alltoallv。

候选词怎么查

主人拿到自己的拼错的词之后，就会生成候选，然后去查它们在不在词典里。不过这时候候选的主人就不一定是自己了，如果在本地，直接查哈希表就行，如果在别的进程，就得发请求。

这里的请求数量其实很吓人。在 64 个进程的时候，三千六百多万次候选里，几乎全都属于别人。可是这些候选绝大部分压根不是词，白白发过去，又白白回来。

于是我给每个进程加了一个 Bloom filter。每个进程把自己那份词典放进一个 Bloom filter 里，然后用 MPI_Allgather 把所有人的过滤器交换一遍，这样每个进程手里都有所有人的过滤器。发请求之前先在本地问一下，如果过滤器说不在，那就肯定不在，这个候选也就不用发了。

for (const std::string& cand : candidates) {
  int owner = GetOwnerRank(cand, world_size);
  if (owner == rank) {
    if (local_dict.Contains(cand)) verified.insert(cand);   // 在本地，直接查
  } else if (MayContain(bloom_of[owner], cand)) {
    remote[owner].push_back(cand);                          // 过滤器说可能在，才发出去
  }                                                         // 说不在的，直接丢掉
}

for (int owner = 0; owner < world_size; ++owner) {
  if (remote[owner].empty()) continue;
  SendCandidateBatch(owner, word_idx, remote[owner]);
  for (int i : WaitForResponse(owner, local_dict))
    verified.insert(remote[owner][i]);
}

每个进程的过滤器有词典总词数的 12 倍那么多位，也就是 750 万位，大概 940KB，用 3 个哈希函数。到了 64 个进程的时候，一个进程要存 64 份，加起来大概 60MB。

效果很明显，三千六百多万次候选里，有约 99.7% 在本地就被挡住了，最后只有 93897 个真的发了出去。我把 Bloom filter 关掉对比了一下，64 个进程的时候，候选阶段要 0.93 秒，打开之后是 0.31 秒。

整个流程连起来，大概是这样的。

实验怎么做

参考实现自带一个简单的测试，可是要看性能的话还不够，所以我自己写了一个 Python 脚本 batch_srun.py 来批量跑。

程序本身会在每个阶段记下自己花的时间，再用 MPI_Reduce 汇总出平均、最小和最大，打印到标准输出里，比如单进程的时候是这样。

[MPI][timing] dict_build avg=0.339734s min=0.339734s max=0.339734s
[MPI][timing] candidate_phase avg=19.1703s min=19.1703s max=19.1703s
[MPI][timing] total_mpi avg=19.9893s min=19.9893s max=19.9893s

脚本负责用 srun 启动程序，然后拿正则把这些数字抠出来，一行一行写进 CSV。

cmd = ["srun", "-n", str(parallel), "--cpus-per-task=1", "build/spellcheck", dict_path, words_path]
proc = subprocess.run(cmd, text=True, capture_output=True)

RE_TIMING = re.compile(r"^\[MPI\]\[timing\]\s+([a-zA-Z0-9_]+)\s+avg=([0-9.]+)s\s+min=([0-9.]+)s\s+max=([0-9.]+)s")

每跑完一次，脚本还会拿输出文件的 MD5 去和事先记下的正确结果比，对不上就直接报错。这样一来，表里每一个速度的数据，都是在结果正确的前提下测出来的。

数据一共有五组，词典和词表分别是 1 万或者 10 万个词的各种搭配，再加上最大的那一组。每一组都从 1 个进程跑到 64 个进程，一共 320 次，所有的进程都在 Spartan 的同一个 sapphire 节点上。这里的数据按 MPI 进程划分，没有测试跨计算节点网络下的扩展性。

不过每个点我只跑了一次，所以曲线上难免会有一些毛刺。比如 12 个进程的那一次，就比旁边的 11 个和 13 个都慢了一截。

结果

下面是最大的那组数据，词典 62 万个词，词表 65 万个词，打开了 Bloom filter。总耗时由外部脚本测量，阶段耗时采用程序输出的 total_mpi。两者的差额还包含启动及未纳入这一计时区间的 I/O 等开销。

候选阶段扩展得非常好，从 19 秒降到 0.31 秒，64 个进程大概快了 62 倍。但是整个程序并没有这么快，64 个进程的总耗时是 1.88 秒，只比单进程快了 10.9 倍。

问题出在别的地方。候选阶段变得很短之后，原来不起眼的那些部分，一下子就显得很大了。

• 词表每个进程都要自己解析一遍，这一步一直是 0.09 秒左右，进程再多也不会变少。
• 建词典是先读文件，再把词发给主人。进程数从 1 加到 8 的时候它确实在变快，从 0.34 秒到 0.09 秒，可是再往上就开始涨了，到 64 个进程的时候是 0.51 秒。毕竟进程多了，每个人读的那一块就变小，而消息的数量却是按进程数的平方在增长。
• 阶段计时之外的启动和 I/O 等开销，也体现在总耗时比阶段耗时多出来的约 0.8 秒里。

所以到了 64 个进程的时候，候选阶段只占阶段耗时的三成不到，建词典反而占了将近一半。

曲线的样子也就好解释了。大概到 32 个进程之后，曲线基本就是平的，总耗时在 1.3 到 2 秒之间晃，最快的一次是 38 个进程的 1.35 秒，大概快了 15 倍。再往后到 64 个进程，反而比 32 个进程略慢。

转发请求中实际命中的比例也跟着进程数在变。因为每个进程的过滤器大小是固定的，进程少的时候，每个人手里的分片大，过滤器就比较满，于是误报也多。2 个进程的时候，被转发出去的候选里只有 19% 是真的在词典里，8 个进程是 88%，16 个进程是 98%，到了 64 个进程就变成 99.96%。

最后和最开始的 BK-tree 比一下，32 分钟变成了两秒以内，其中一百倍来自换算法，剩下的十到十五倍来自并行。

报告

作为作业的一部分，最终写了一份完整的报告，包含实验设置、结果与分析。下面可以直接预览或下载原始 PDF。