计算144维上的欧几里得距离

去年下半年，我读了一篇有关CSAM图像扫描工具的博客文章。我记得当时在想：这太酷了！图像处理总是很困难，在Cloudflare上部署真实的图像识别系统可不是一件容易的事！

一段时间后，我和Kornel聊天：＆＃34;我们在图像处理管道中拥有所有组件，但是我们在为一个组件的性能而苦苦挣扎。扩展到Cloudflare并非易事！

问题在于匹配算法本身的速度。让我详细说明。如John在其博客文章中所述，图像匹配算法会根据处理后的图像创建模糊哈希。哈希正好是144个字节长。例如，它可能看起来像这样：

哈希设计用于模糊匹配算法中，该算法可以找到附近的相关图像。特定的算法定义明确，但是要使其快速完成则由程序员负责–在Cloudflare，我们需要以超快的速度完成匹配。我们希望将每秒数以千计的散列通过我们的网络传递的图像与数百万个已知图像的数据库进行匹配。为此，我们需要认真优化匹配算法。

我想到的第一个算法具有O（K * N）复杂度：对于每个查询，都要遍历数据库中的每个哈希。在幼稚的实现中，这会产生很多工作。但是到底有多少工作？

给定查询散列，模糊匹配为＆＃34; closest＆＃34;。在数据库中哈希。这要求我们定义一个距离。我们将每个哈希视为包含144个数字的向量，以标识144维空间中的一个点。给定两个这样的点，我们可以使用标准的欧几里得公式来计算距离。

不过，对于我们的特定问题，我们对＆＃34; closest＆＃34;感兴趣。仅当距离小于某个预定义阈值时才匹配数据库。否则，当距离较大时，我们可以假定图像不相似。这是预期的结果-我们的大多数查询在数据库中都没有相关的图像。

要计算两个144字节散列之间的距离，我们取每个字节，计算增量，将其平方，然后将其求和到一个累加器中，求平方根，然后ta-dah！我们有距离！

此函数返回距离的平方。我们避免计算实际距离，以免我们无法运行平方根函数-它很慢。在代码内部，为了提高性能和简化操作，我们主要对平方值进行操作。我们不需要实际的距离值，只需要找到最小的向量即可。对于我们来说，比较距离还是平方距离并不重要！

如您所见，模糊匹配基本上是在多维空间中找到最接近点的标准问题。当然，过去已经解决了这个问题-但我们不要跳过。

虽然这段代码可能很简单，但我们希望它会很慢。要在数据库中找到最小的哈希距离（例如1M条目），将需要遍历所有记录，并且至少需要：

和更多。仅此一项，就增加了4.32亿次操作！实际情况如何？为了说明此博客文章，我们准备了完整的测试套件。大型的已知哈希数据库可以通过随机数据很好地模拟。查询散列不能是随机的，并且必须稍微复杂一些，否则，练习就不会那么有趣了。我们通过字节交换数据库中的实际数据来智能地生成测试-这使我们能够精确地控制测试哈希与数据库哈希之间的距离。查看脚本以了解详细信息。这是我们首次运行的第一个幼稚算法：

$天真＆lt; test-vector.txt ./mmdist-naive＆gt; test-vector.tmp总计：85261.833ms，1536项，平均每次查询55.509ms，18.015 qps

我们在85秒内将1536个测试哈希与100万个随机向量的数据库进行了匹配。找到最接近的邻居平均花了55ms CPU时间。这对于我们的需求而言相当慢。

一个明显的改进是使用了更复杂的SIMD指令。 SIMD是一种指示CPU使用一条指令处理多个数据点的方法。当处理向量问题时，这是一个完美的策略-正如我们的任务一样。

我们决定使用具有256位向量的AVX2。我们这样做的原因很简单-AMD CPU不支持较新的AVX版本。此外，过去，我们对AVX-512的频率缩放并不感到兴奋。

说使用AVX2容易做起来难。没有任何一条指令可以计算两个uint8向量之间的欧几里得距离！由Vlad编写的用AVX2计算两个144字节向量的完整距离的最快方法是由Vlad编写的：

它实际上比看起来简单：加载16个字节，将向量从uint8转换为int16，减去向量，将中间和存储为int32，然后重复。最后，我们需要执行4条复杂的指令以将部分和提取为最终和。此AVX2代码将性能提高了3倍左右：

我们测量的每个项目为17毫秒，仍然低于我们的预期。不幸的是，如果没有重大更改，我们将无法进一步推进。问题在于此代码受内存带宽限制。测量来自我的Intel i7-5557U CPU，其最大理论内存带宽仅为25GB / s。 100万个条目的数据库需要137MiB，因此至少需要5毫秒才能将数据库提供给我的CPU。使用这种幼稚的算法，我们将无法做到这一点。

由于朴素的暴力破解方法失败，因此我们尝试使用更复杂的算法。我的同事KornelLesiński实现了超酷的Vantage Point算法。经过几次跌宕起伏，优化和重写之后，我们放弃了。对于这种算法，我们的问题变得异常困难。

我们观察到了“维度的诅咒”。空间划分算法在处理大维问题时效果不佳-在我们的案例中，我们有144个维。 K-D树注定要失败。局域性哈希也注定要失败。这是一个奇怪的情况，其中的空间是难以想象的巨大，但是所有东西都紧密地结合在一起。该空间的大小为347位数字长，但点之间的最大距离仅为3060-sqrt（255 * 255 * 144）。

空间划分算法之所以快速，是因为随着它们越来越靠近寻找最近的点，它们逐渐缩小了搜索空间。但是在我们的情况下，通用查询永远不会接近集合中的任何一点，因此搜索空间无法缩小到有意义的程度。

VP树是一个很有前途的候选者，因为它仅在一定距离上运行，将空间细分为近和远的分区，就像二叉树一样。当它具有紧密匹配时，它可以非常快，并且不需要访问超过O（log（N））个节点。对于不匹配的游戏，其速度会急剧下降。该算法最终访问树中几乎一半的节点。一切都在144个维度上紧密结合在一起！即使该算法避免访问树中超过一半的节点，但访问其余节点的成本较高，因此搜索最终总体上变慢了。

这次经历使我们开始思考。由于空间划分算法无法缩小搜索范围，并且仍然需要遍历大量项目，因此也许我们应该专注于非常快速地遍历所有哈希。但是，我们必须对内存带宽更加机敏-以前它是天真的暴力方法的限制因素。

突破来自认识到我们不需要计算散列之间的完整距离。相反，我们只能计算尺寸的一个子集，例如144个尺寸中的32个。如果该距离已经很大，则无需计算整个尺寸！计算更多的点不会减少欧几里得距离。

2.遍历数据库中所有的100万个32字节短哈希。它们必须密集地包装在内存中，以使CPU能够执行良好的预取，并避免读取我们不需要的数据。

3.如果到目前为止32字节短哈希的距离大于或等于最佳分数，请继续

即使此算法需要执行较少的算术和内存工作，也不会比以前的幼稚算法快。请参见make short-avx2。问题是：我们仍然需要为有前途的哈希计算完整距离，并且其中有很多。计算有希望的哈希的完整距离会增加ALU和内存延迟方面的工作量，以抵消该算法的收益。

我们对图像匹配问题的特殊应用有一个细节，它将帮助我们取得更大的进步。如我们先前所述，问题不在于寻找最接近的邻居，而在于证明不存在具有合理距离的邻居。请记住-实际上，我们不希望找到很多比赛！我们期望我们输入到算法中的几乎所有图像都与存储在数据库中的图像哈希无关。

对于我们的算法而言，足以证明在预定的距离阈值内不存在邻居。假设我们对比220距离更远的哈希值不感兴趣，例如220,400的平方。这使我们的短距离算法变体效果更好：

在某个时候，John指出该阈值允许进行其他优化。我们可以通过哈希值到某个原点的距离对其进行排序。给定查询散列的起始距离为A，我们只能检查| A-threshold |之间的散列和| A +阈值|从起源。这就是Vantage点树的每个级别的工作原理，只是进行了简化。这种优化-按距原点的距离对数据库中的项目进行排序-比较简单，可以帮助我们节省一些工作。

尽管在纸面上很不错，但是这种方法在实践中并没有带来太大的收益，因为矢量没有按聚类分组-它们几乎是随机的！对于我们感兴趣的阈值，原点距离算法的变化使我们的速度提高了约20％，这虽然可以，但令人叹为观止。如果我们决定降低阈值，此更改可能会带来更多好处，因此对于生产实施可能值得这样做。但是，它不适用于查询批处理。

但是我们还没有完成AVX优化！ AVX的常见问题是说明通常不适合特定问题。为了使正确的指令适应问题或使问题逆转，需要认真思考，以便可以使用特定的指令。 AVX2没有有用的＆＃34;水平＆＃34; uint16减，乘和加运算。例如，存在_mm_hadd_epi16，但是它又慢又麻烦。

相反，我们可以解决问题以利用快速可用的uint16操作数。例如，我们可以使用：

饱和添加非常有用，可以节省我们转换为uint32的时间。它只是增加了一个小限制：传递给程序的阈值（即最大平方距离）必须适合uint16。但是，这对我们很好。

为了有效地使用这些指令，我们需要转置数据库中的数据。除了将哈希存储在行中之外，我们还可以将它们存储在列中：

现在，我们可以使用一个内存操作加载16个哈希的第一个字节。在下一步中，我们可以使用一条指令减去查询哈希的第一个字节，依此类推。该算法与上面定义的完全相同；我们只是使数据更易于加载和更易于AVX处理。

通过调整好的批次大小和短距离大小参数，我们可以看到此算法的性能：

哇！太棒了我们从每个查询55毫秒开始，仅以0.73毫秒结束。还有可能进行进一步的微优化，例如内存预取或使用大页面来减少页面错误，但此时它们的收益递减。

如果您对这样的架构调整感兴趣，请阅读Denis Bakhvalov撰写的新性能手册。它讨论了车顶线模型分析，这几乎就是我们在这里所做的。

请看一下我们的代码，并告诉我们是否错过了一些优化！

多么优化的旅程！我们在内存和ALU瓶颈代码之间跳转。我们讨论了更复杂的算法，但最后，尽管进行了调整，但蛮力算法给了我们最好的结果。

为了获得更好的数字，我使用CUDA对Nvidia GPU进行了实验。像vabsdiff4和dp4a这样的CUDA内在函数完全适合该问题。 V100给了我们一些惊人的数字，但我并不完全满意。考虑到我们可以用一个服务器级GPU的成本获得多少个带有AVX2的AMD Ryzen内核，我们针对这个特殊问题倾向于通用计算。

这是我们每天处理的复杂类型的一个很好的例子。要使最佳技术“在Cloudflare规模上”发挥作用，就需要跳出思路。有时，在找到最佳解决方案之前，我们重写了数十次解决方案。有时我们会采用蛮力算法，只是非常非常优化。

散列的计算和图像匹配是具有挑战性的问题，需要运行非常占用CPU的操作。我们在边缘上可用的CPU稀缺，这样的工作量非常昂贵。即使在本博客文章中讨论了优化工作，大规模运行CSAM扫描器也是一个挑战，并且需要大量的工程工作。我们还没有完成！在我们满意之前，我们需要解决更多的难题。如果您想提供帮助，请考虑申请！

性能优化速度