Rust中的快速慣用Floyd-Warshall演算法

我正在嘗試在 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 倍的速度運行（測驗的完整代碼）。

我沒想到邊界檢查會增加那么多開銷:(

是否可以在沒有 unsafes 的情況下以慣用的方式重寫此代碼，使其與 unsafe 版本一樣快？例如，是否可以通過向代碼添加一些斷言來向編譯器“證明”不會有越界訪問？

uj5u.com熱心網友回復：

乍一看，人們希望這足夠了：

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 的建議使用了 2D 陣列：

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 毫秒。在檢查程式集時（#[inline(never)]在 floyd_warshall 上），兩者似乎都沒有進行邊界檢查，而且兩者都在某種程度上看起來是矢量化的。雖然，我不是閱讀匯編的專家。

因為這是一個如此熱的回圈（最多三個邊界檢查），我對性能受到如此大的影響并不感到驚訝。不幸的是，在這種情況下索引的使用非常復雜，以至于斷言技巧無法為您提供簡單的解決方案。還有其他已知的情況，其中需要斷言檢查以提高性能，但編譯器無法充分使用資訊。這是一個這樣的例子。

這是我改變的游樂場。

uj5u.com熱心網友回復：

經過一些實驗，基于安德魯回答中建議的想法以及相關問題中的評論，我找到了解決方案，其中：

• 仍然使用相同的介面（例如&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]）來完成的。之后，編譯器知道內回圈永遠不會使用越界值，并洗掉檢查。

有趣的是，只有當所有三個優化都使用時，它才能大大提高速度。如果我們只使用其中的任意兩個，編譯器將無法對其進行優化。

標籤：算法 锈 图算法

上一篇：如何創建另一個保證預測的陣列集

下一篇：從IEnumerable<T>函式回傳字串