Rust 中的快速惯用 Floyd-Warshall 算法

Posted 2023-02-19

【中文标题】Rust 中的快速惯用 Floyd-Warshall 算法

【英文标题】：Fast idiomatic Floyd-Warshall algorithm in Rust

【发布时间】：2021-12-31 03:00:12

【问题描述】：

我正在尝试在 Rust 中实现一个相当快的 Floyd-Warshall 算法版本。该算法在有向加权图中找到所有顶点之间的最短路径。

算法的主要部分可以这样写：

// dist[i][j] contains edge length between vertices [i] and [j]
// after the end of the execution it contains shortest path between [i] and [j]
fn floyd_warshall(dist: &mut [Vec<i32>]) 
    let n = dist.len();
    for i in 0..n 
        for j in 0..n 
            for k in 0..n 
                dist[j][k] = min(dist[j][k], dist[j][i] + dist[i][k]);
            
        
    

此实现非常简短且易于理解，但它的运行速度比类似的 c++ 实现慢 1.5 倍。

据我了解，问题是在每个向量访问时，Rust 都会检查索引是否在向量的边界内，这会增加一些开销。

我用 get_unchecked* 函数重写了这个函数：

fn floyd_warshall_unsafe(dist: &mut [Vec<i32>]) 
    let n = dist.len();
    for i in 0..n 
        for j in 0..n 
            for k in 0..n 
                unsafe 
                    *dist[j].get_unchecked_mut(k) = min(
                        *dist[j].get_unchecked(k),
                        dist[j].get_unchecked(i) + dist[i].get_unchecked(k),
                    )
                
            
        
    

它的运行速度确实提高了 1.5 倍 (full code of the test)。

我没想到边界检查会增加这么多开销:(

是否可以在没有不安全的情况下以惯用的方式重写此代码，使其与不安全版本一样快？例如。是否可以通过在代码中添加一些断言来向编译器“证明”不会有越界访问？

【问题讨论】：

您是否与 Vecs 数组（或其他任何东西）结婚？我的第一个想法是切换到一个合适的二维数组，或者失败，一个手动索引的一维数组。然后，您可以通过断言一维数组的长度为 n*n 来说服编译器放弃边界检查。

另外，你是用--release编译的，对吧？

是的，@Jmb，我正在发布模式下编译。

至于@DavidEisenstat 点-在Rust 世界中什么被认为是合适的二维数组？我尝试了array2d crate，但它的运行速度甚至比 vecs 的 Vec 还要慢。我还基于code 内部的一维向量实现了我自己的 Array2D，它的工作速度比不安全版本慢约 10%（这是我对每个向量访问的一次边界检查所期望的），但它比 Vecs 的 Vec 好得多版本！

我不是 Rust 程序员，所以我不知道。在幕后，LLVM 似乎不理解 2D 数组，而且这个 C++ 测试程序也没有像希望的那样优化，所以我对回答这个问题的前景感到悲观：#include <cassert> void test(int n) assert(n >= 0); for (int i = 0; i < n; i++) for (int j = 0; j < n; j++) assert(i + j < n + n); 。

【参考方案1】：

乍一看，希望这已经足够了：

fn floyd_warshall(dist: &mut [Vec<i32>]) 
    let n = dist.len();
    for i in 0..n 
        assert!(i < dist.len());
        for j in 0..n 
            assert!(j < dist.len());
            assert!(i < dist[j].len());
            let v2 = dist[j][i];
            for k in 0..n 
                assert!(k < dist[i].len());
                assert!(k < dist[j].len());
                dist[j][k] = min(dist[j][k], v2 + dist[i][k]);
            
        
    

添加断言是让 Rust 优化器相信变量确实在界限内的已知技巧。但是，它在这里不起作用。我们需要做的是以某种方式让 Rust 编译器更清楚地知道这些循环是有界的，而无需求助于深奥的代码。

为了实现这一点，我按照 David Eisenstat 的建议转移到了二维数组：

fn floyd_warshall<const N:usize>(mut dist: Box<[[i32; N]; N]>) -> Box<[[i32; N]; N]> 
    for i in 0..N 
        for j in 0..N 
            for k in 0..N 
                dist[j][k] = min(dist[j][k], dist[j][i] + dist[i][k]);
            
        
    
    dist

这使用常量泛型（Rust 的一个相对较新的特性）来指定堆上给定二维数组的大小。就其本身而言，此更改在我的机器上运行良好（比 usafe 快 100 毫秒，比 unsafe 慢约 20 毫秒）。另外，如果您将 v2 计算移到 k 循环之外，如下所示：

fn floyd_warshall<const N:usize>(mut dist: Box<[[i32; N]; N]>) -> Box<[[i32; N]; N]> 
    for i in 0..N 
        for j in 0..N 
            let v2 = dist[j][i];
            for k in 0..N 
                dist[j][k] = min(dist[j][k], v2 + dist[i][k]);
            
        
    
    dist

改进是巨大的（在我的机器上从 ~300ms 到 ~100ms）。同样的优化适用于floyd_warshall_unsafe，在我的机器上平均达到约 100 毫秒。在检查程序集时（在 floyd_warshall 上使用 #[inline(never)]），看起来两者都没有发生边界检查，并且两者看起来都在某种程度上是矢量化的。虽然，我不是阅读汇编的专家。

因为这是一个非常热的循环（最多进行三个边界检查），所以性能受到如此大的影响并不奇怪。不幸的是，在这种情况下，索引的使用非常复杂，以至于断言技巧无法为您提供简单的修复。还有其他一些已知情况，需要进行断言检查以提高性能，但编译器无法充分使用这些信息。 Here is one such example.

Here is the playground 我的更改。

【讨论】：

哇，这个let v2 = dist[j][i] 真的让它更快！可能编译器意识到我们对所有k 都做了类似的事情，并且能够向量化代码。

我同意常量泛型有助于消除边界检查，但这似乎有点不切实际，因为如果我们在编译时不知道图形的大小，它就不起作用，对吧？跨度>

请注意，断言是错误的，因为它们根据外部向量 (dist) 的长度而不是内部向量的长度检查所有三个变量，但没有明显的原因为什么它们应该是一样的。所以你应该断言i < dist.len()、j < dist.len()、i < dist[j].len()、k < dist[i].len()和k < dist[j].len()。

好点@Jmb，我修复了代码并再次测试它，不幸的是它仍然没有忽略边界检查。

@Borys 确实需要在编译时知道数组的大小。在这方面，不安全的代码更实用。如果您真的想使用 const 通用代码，则需要选择具体的尺寸并将您的图形调整到该尺寸以供使用。您可能要考虑的一种方法是使用 Prusti 为您未在边界之外访问的语义 cmets 添加合同

【参考方案2】：

经过一些实验，基于Andrew's answer和comments in related issue中提出的想法，我找到了解决方案：

仍然使用相同的接口（例如 &mut [Vec<i32>] 作为参数） 不使用不安全 比不安全版本快 3-4 倍 很丑:(

代码如下：

fn floyd_warshall_fast(dist: &mut [Vec<i32>]) 
    let n = dist.len();
    for i in 0..n 
        for j in 0..n 
            if i == j 
                continue;
            
            let (dist_j, dist_i) = if j < i 
                let (lo, hi) = dist.split_at_mut(i);
                (&mut lo[j][..n], &mut hi[0][..n])
             else 
                let (lo, hi) = dist.split_at_mut(j);
                (&mut hi[0][..n], &mut lo[i][..n])
            ;
            let dist_ji = dist_j[i];
            for k in 0..n 
                dist_j[k] = min(dist_j[k], dist_ji + dist_i[k]);
            
        
    

里面有几个想法：

我们计算一次dist_ji，因为它在最内循环内不会改变，编译器不需要考虑它。 我们“证明”dist[i] 和 dist[j] 实际上是两个不同的向量。这是由这个丑陋的split_at_mut 和i == j 特殊情况完成的（真的很想知道一个更简单的解决方案）。之后我们可以将dist[i] 和dist[j] 完全分开处理，例如编译器可以向量化这个循环，因为它知道数据不会重叠。 最后一个技巧是向编译器“证明”dist[i] 和dist[j] 至少具有n 元素。这是由[..n] 在计算dist[i] 和dist[j] 时完成的（例如，我们使用&mut lo[j][..n] 而不仅仅是&mut lo[j]）。之后，编译器了解内部循环永远不会使用超出范围的值，并删除检查。

有趣的是，只有当所有三个优化都使用时，它才能大大加快速度。如果我们只使用其中任何两个，编译器无法优化它。

【讨论】：

代码似乎是用外来脚本编写的：P，问问自己优化是否值得不可读。事后猜测或强制编译器进行优化是一种已知的反模式，这会导致脆弱的代码随着编译器的发展而中断。