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

Question

我正在尝试在 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)。

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

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

Answer 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 我的更改。

Answer 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]）。之后，编译器了解内部循环永远不会使用超出范围的值，并删除检查。

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