在 Haskell 中执行常量空间嵌套循环的正确方法是什么？

Question

在 Haskell 中执行嵌套循环有两种明显的“惯用”方式：使用 list monad 或使用 forM_ 替换传统的 fors。我已经设置了一个基准来确定这些是否被编译为紧密循环：

import Control.Monad.Loop
import Control.Monad.Primitive
import Control.Monad
import Control.Monad.IO.Class
import qualified Data.Vector.Unboxed.Mutable as MV
import qualified Data.Vector.Unboxed as V

times = 100000
side  = 100

-- Using `forM_` to replace traditional fors
test_a mvec = 
    forM_ [0..times-1] $ \ n -> do
        forM_ [0..side-1] $ \ y -> do
            forM_ [0..side-1] $ \ x -> do
                MV.write mvec (y*side+x) 1

-- Using the list monad to replace traditional forms
test_b mvec = sequence_ $ do
    n <- [0..times-1]
    y <- [0..side-1]
    x <- [0..side-1]
    return $ MV.write mvec (y*side+x) 1

main = do
    let vec = V.generate (side*side) (const 0)
    mvec <- V.unsafeThaw vec :: IO (MV.MVector (PrimState IO) Int)
    -- test_a mvec
    -- test_b mvec
    vec' <- V.unsafeFreeze mvec :: IO (V.Vector Int)
    print $ V.sum vec'

此测试创建一个 100x100 的向量，使用嵌套循环将 1 写入每个索引并重复 100k 次。仅使用 ghc -O2 test.hs -o test（ghc 版本 7.8.4）编译那些，结果是：forM_ 版本为 3.853s，@987654329 版本为 10.460s @。为了提供参考，我还用 JavaScript 编写了这个测试：

var side  = 100;
var times = 100000;
var vec   = [];

for (var i=0; i<side*side; ++i)
    vec.push(0);

for (var n=0; n<times; ++n)
    for (var y=0; y<side; ++y)
        for (var x=0; x<side; ++x)
            vec[x+y*side] = 1;

var s = 0;
for (var i=0; i<side*side; ++i)
    s += vec[i];

console.log(s);

这个等效的 JavaScript 程序需要 1s 来完成，击败了 Haskell 的未装箱向量，这是不寻常的，这表明 Haskell 不是在恒定空间中运行循环，而是进行分配。然后我找到了一个声称提供类型保证紧密循环的库Control.Monad.Loop：

-- Using `for` from Control.Monad.Loop
test_c mvec = exec_ $ do
    n <- for 0 (< times) (+ 1)
    x <- for 0 (< side) (+ 1)
    y <- for 0 (< side) (+ 1)
    liftIO (MV.write mvec (y*side+x) 1)

在 1s 中运行。但是，该库不是很常用，而且远非惯用，所以，获得快速恒定空间二维计算的惯用方法是什么？（请注意，这不是我想要的 REPA 的情况在网格上执行任意 IO 操作。）

Answer 1

使用 GHC 编写严格的变异代码有时会很棘手。我将写一些不同的事情，可能以一种比我更喜欢的方式更漫无边际和 tl;dr。

对于初学者，无论如何我们都应该使用 GHC 7.10，因为 otherwise forM_ 和 list monad 解决方案永远不会融合。

另外，我用MV.unsafeWrite 替换了MV.write，部分原因是它更快，但更重要的是它减少了生成的Core 中的一些混乱。从现在开始，运行时统计信息请参考带有unsafeWrite 的代码。

可怕的让漂浮

即使使用 GHC 7.10，我们也应该首先注意到所有这些 [0..times-1] 和 [0..side-1] 表达式，因为如果我们不采取必要的步骤，它们每次都会破坏性能。问题是它们是恒定范围，-ffull-laziness（默认情况下在-O 上启用）将它们浮动到顶层。这可以防止列表融合，并且迭代 Int# 范围比迭代盒装 Int-s 的列表便宜，所以这是一个非常糟糕的优化。

让我们在几秒钟内查看未更改（除了使用unsafeWrite）代码的一些运行时。使用ghc -O2 -fllvm，我使用+RTS -s进行计时。

test_a: 1.6
test_b: 6.2
test_c: 0.6

对于 GHC Core 查看，我使用了ghc -O2 -ddump-simpl -dsuppress-all -dno-suppress-type-signatures。

在test_a 的情况下，[0..99] 范围被提升：

main4 :: [Int]
main4 = eftInt 0 99 -- means "enumFromTo" for Int.

虽然最外层的[0..9999] 循环被融合成一个尾递归助手：

letrec {
          a3_s7xL :: Int# -> State# RealWorld -> (# State# RealWorld, () #)
          a3_s7xL =
            \ (x_X5zl :: Int#) (s1_X4QY :: State# RealWorld) ->
              case a2_s7xF 0 s1_X4QY of _ { (# ipv2_a4NA, ipv3_a4NB #) ->
              case x_X5zl of wild_X1S {
                __DEFAULT -> a3_s7xL (+# wild_X1S 1) ipv2_a4NA;
                99999 -> (# ipv2_a4NA, () #)
              }
              }; }

对于test_b，同样只有[0..99] 被解除。但是，test_b 慢得多，因为它必须构建和排序实际的 [IO ()] 列表。至少 GHC 足够明智，只为两个内循环构建一个[IO ()]，然后对其执行10000 次排序。

 let {
          lvl7_s4M5 :: [IO ()]
          lvl7_s4M5 = -- omitted
        letrec {
          a2_s7Av :: Int# -> State# RealWorld -> (# State# RealWorld, () #)
          a2_s7Av =
            \ (x_a5xi :: Int#) (eta_B1 :: State# RealWorld) ->
              letrec {
                a3_s7Au
                  :: [IO ()] -> State# RealWorld -> (# State# RealWorld, () #)
                a3_s7Au =
                  \ (ds_a4Nu :: [IO ()]) (eta1_X1c :: State# RealWorld) ->
                    case ds_a4Nu of _ {
                      [] ->
                        case x_a5xi of wild1_X1y {
                          __DEFAULT -> a2_s7Av (+# wild1_X1y 1) eta1_X1c;
                          99999 -> (# eta1_X1c, () #)
                        };
                      : y_a4Nz ys_a4NA ->
                        case (y_a4Nz `cast` ...) eta1_X1c
                        of _ { (# ipv2_a4Nf, ipv3_a4Ng #) ->
                        a3_s7Au ys_a4NA ipv2_a4Nf
                        }
                    }; } in
              a3_s7Au lvl7_s4M5 eta_B1; } in
-- omitted

我们如何解决这个问题？我们可以用{-# OPTIONS_GHC -fno-full-laziness #-} 解决这个问题。这在我们的案例中确实有很大帮助：

test_a: 0.5
test_b: 0.48
test_c: 0.5

或者，我们可以摆弄INLINE pragmas。显然，在 let 浮动完成后内联函数可以保持良好的性能。我发现即使没有 pragma，GHC 也会内联我们的测试函数，但是显式的 pragma 会导致它仅在 let 浮动后才内联。例如，这会导致没有-fno-full-laziness 的良好性能：

test_a mvec = 
    forM_ [0..times-1] $ \ n -> 
        forM_ [0..side-1] $ \ y -> 
            forM_ [0..side-1] $ \ x -> 
                MV.unsafeWrite mvec (y*side+x) 1
{-# INLINE test_a #-}

但过早内联会导致性能不佳：

test_a mvec = 
    forM_ [0..times-1] $ \ n -> 
        forM_ [0..side-1] $ \ y -> 
            forM_ [0..side-1] $ \ x -> 
                MV.unsafeWrite mvec (y*side+x) 1
{-# INLINE [~2] test_a #-} -- "inline before the first phase please"

这个INLINE 解决方案的问题在于，面对 GHC 的浮动冲击，它相当脆弱。例如，手动内联不会保持性能。以下代码很慢，因为与INLINE [~2] 类似，它让 GHC 有机会浮出水面：

main = do
    let vec = V.generate (side*side) (const 0)
    mvec <- V.unsafeThaw vec :: IO (MV.MVector (PrimState IO) Int)
    forM_ [0..times-1] $ \ n -> 
        forM_ [0..side-1] $ \ y -> 
            forM_ [0..side-1] $ \ x -> 
                MV.unsafeWrite mvec (y*side+x) 1    

那么我们应该怎么做呢？

首先，我认为使用-fno-full-laziness 对于那些想要编写高性能代码并清楚自己在做什么的人来说是一个完全可行甚至更可取的选择。例如，它用于unordered-containers。有了它，我们可以更精确地控制共享，我们总是可以手动浮动或内联。

对于更常规的代码，我相信使用Control.Monad.Loop 或任何其他提供该功能的包没有任何问题。许多 Haskell 用户对依赖小型“边缘”库并不谨慎。我们也可以只重新实现for，以达到期望的通用性。例如，以下解决方案的性能与其他解决方案一样好：

for :: Monad m => a -> (a -> Bool) -> (a -> a) -> (a -> m ()) -> m ()
for init while step body = go init where
  go !i | while i = body i >> go (step i)
  go i = return ()
{-# INLINE for #-}

在真正恒定的空间中循环

起初我对+RTS -s 堆分配的数据感到非常困惑。 test_a 与-fno-full-laziness 分配非平凡，并且test_c 没有完全惰性，这些分配与times 迭代的数量成线性比例，但test_b 分配了完全惰性仅适用于矢量：

-- with -fno-full-laziness, no INLINE pragmas
test_a: 242,521,008 bytes
test_b: 121,008 bytes
test_c: 121,008 bytes -- but 240,120,984 with full laziness!

此外，test_c 的 INLINE 编译指示在这种情况下根本没有帮助。

我花了一些时间试图在Core中为相关程序找到堆分配的迹象，但没有成功，直到我意识到：GHC堆栈帧在堆上，包括主线程的帧，以及函数进行堆分配的那些基本上是在最多三个堆栈帧中运行三次嵌套循环。 +RTS -s注册的堆分配只是栈帧的不断弹出和推送。

以下代码的核心非常明显：

{-# OPTIONS_GHC -fno-full-laziness #-}

-- ...

test_a mvec = 
    forM_ [0..times-1] $ \ n -> 
        forM_ [0..side-1] $ \ y -> 
            forM_ [0..side-1] $ \ x -> 
                MV.unsafeWrite mvec (y*side+x) 1
main = do
    let vec = V.generate (side*side) (const 0)
    mvec <- V.unsafeThaw vec :: IO (MV.MVector (PrimState IO) Int)
    test_a mvec

我将其包括在它的荣耀中。随意跳过。

main1 :: State# RealWorld -> (# State# RealWorld, () #)
main1 =
  \ (s_a5HK :: State# RealWorld) ->
    case divInt# 9223372036854775807 8 of ww4_a5vr { __DEFAULT ->

    -- start of vector creation ----------------------
    case tagToEnum# (># 10000 ww4_a5vr) of _ {
      False ->
        case newByteArray# 80000 (s_a5HK `cast` ...)
        of _ { (# ipv_a5fv, ipv1_a5fw #) ->
        letrec {
          $s$wa_s8jS
            :: Int#
               -> Int#
               -> State# (PrimState IO)
               -> (# State# (PrimState IO), Int #)
          $s$wa_s8jS =
            \ (sc_s8jO :: Int#)
              (sc1_s8jP :: Int#)
              (sc2_s8jR :: State# (PrimState IO)) ->
              case tagToEnum# (<# sc1_s8jP 10000) of _ {
                False -> (# sc2_s8jR, I# sc_s8jO #);
                True ->
                  case writeIntArray# ipv1_a5fw sc_s8jO 0 (sc2_s8jR `cast` ...)
                  of s'#_a5Gn { __DEFAULT ->
                  $s$wa_s8jS (+# sc_s8jO 1) (+# sc1_s8jP 1) (s'#_a5Gn `cast` ...)
                  }
              }; } in
        case $s$wa_s8jS 0 0 (ipv_a5fv `cast` ...)
        -- end of vector creation -------------------

        of _ { (# ipv6_a4Hv, ipv7_a4Hw #) ->
        letrec {
          a2_s7MJ :: Int# -> State# RealWorld -> (# State# RealWorld, () #)
          a2_s7MJ =
            \ (x_a5Ho :: Int#) (eta_B1 :: State# RealWorld) ->
              letrec {
                a3_s7ME :: Int# -> State# RealWorld -> (# State# RealWorld, () #)
                a3_s7ME =
                  \ (x1_X5Id :: Int#) (eta1_XR :: State# RealWorld) ->
                    case ipv7_a4Hw of _ { I# dt4_a5x6 ->
                    case writeIntArray#
                           (ipv1_a5fw `cast` ...) (*# x1_X5Id 100) 1 (eta1_XR `cast` ...)
                    of s'#_a5Gn { __DEFAULT ->
                    letrec {
                      a4_s7Mz :: Int# -> State# RealWorld -> (# State# RealWorld, () #)
                      a4_s7Mz =
                        \ (x2_X5J8 :: Int#) (eta2_X1U :: State# RealWorld) ->
                          case writeIntArray#
                                 (ipv1_a5fw `cast` ...)
                                 (+# (*# x1_X5Id 100) x2_X5J8)
                                 1
                                 (eta2_X1U `cast` ...)
                          of s'#1_X5Hf { __DEFAULT ->
                          case x2_X5J8 of wild_X2o {
                            __DEFAULT -> a4_s7Mz (+# wild_X2o 1) (s'#1_X5Hf `cast` ...);
                            99 -> (# s'#1_X5Hf `cast` ..., () #)
                          }
                          }; } in
                    case a4_s7Mz 1 (s'#_a5Gn `cast` ...)
                    of _ { (# ipv2_a4QH, ipv3_a4QI #) ->
                    case x1_X5Id of wild_X1e {
                      __DEFAULT -> a3_s7ME (+# wild_X1e 1) ipv2_a4QH;
                      99 -> (# ipv2_a4QH, () #)
                    }
                    }
                    }
                    }; } in
              case a3_s7ME 0 eta_B1 of _ { (# ipv2_a4QH, ipv3_a4QI #) ->
              case x_a5Ho of wild_X1a {
                __DEFAULT -> a2_s7MJ (+# wild_X1a 1) ipv2_a4QH;
                99999 -> (# ipv2_a4QH, () #)
              }
              }; } in
        a2_s7MJ 0 (ipv6_a4Hv `cast` ...)
        }
        };
      True ->
        case error
               (unpackAppendCString#
                  "Primitive.basicUnsafeNew: length to large: "#
                  (case $wshowSignedInt 0 10000 ([])
                   of _ { (# ww5_a5wm, ww6_a5wn #) ->
                   : ww5_a5wm ww6_a5wn
                   }))
        of wild_00 {
        }
    }
    }

main :: IO ()
main = main1 `cast` ...

main2 :: State# RealWorld -> (# State# RealWorld, () #)
main2 = runMainIO1 (main1 `cast` ...)

main :: IO ()
main = main2 `cast` ...

我们还可以通过以下方式很好地演示帧的分配。换个test_a：

test_a mvec = 
    forM_ [0..times-1] $ \ n -> 
        forM_ [0..side-1] $ \ y -> 
            forM_ [0..side-50] $ \ x -> -- change here
                MV.unsafeWrite mvec (y*side+x) 1

现在堆分配保持不变，因为最里面的循环是尾递归的并且使用单个帧。通过以下更改，堆分配减半（为 124,921,008 字节），因为我们推送和弹出一半的帧：

test_a mvec = 
    forM_ [0..times-1] $ \ n -> 
        forM_ [0..side-50] $ \ y -> -- change here
            forM_ [0..side-1] $ \ x -> 
                MV.unsafeWrite mvec (y*side+x) 1

test_b 和 test_c（没有完全惰性）改为编译为在单个堆栈帧内使用嵌套 case 构造的代码，并遍历索引以查看应该递增的索引。请参阅以下main 的核心：

{-# LANGUAGE BangPatterns #-} -- later I'll talk about this
{-# OPTIONS_GHC -fno-full-laziness #-}

main = do
    let vec = V.generate (side*side) (const 0)
    !mvec <- V.unsafeThaw vec :: IO (MV.MVector (PrimState IO) Int)
    test_c mvec

瞧：

main1 :: State# RealWorld -> (# State# RealWorld, () #)
main1 =
  \ (s_a5Iw :: State# RealWorld) ->
    case divInt# 9223372036854775807 8 of ww4_a5vT { __DEFAULT ->

    -- start of vector creation ----------------------
    case tagToEnum# (># 10000 ww4_a5vT) of _ {
      False ->
        case newByteArray# 80000 (s_a5Iw `cast` ...)
        of _ { (# ipv_a5g3, ipv1_a5g4 #) ->
        letrec {
          $s$wa_s8ji
            :: Int#
               -> Int#
               -> State# (PrimState IO)
               -> (# State# (PrimState IO), Int #)
          $s$wa_s8ji =
            \ (sc_s8je :: Int#)
              (sc1_s8jf :: Int#)
              (sc2_s8jh :: State# (PrimState IO)) ->
              case tagToEnum# (<# sc1_s8jf 10000) of _ {
                False -> (# sc2_s8jh, I# sc_s8je #);
                True ->
                  case writeIntArray# ipv1_a5g4 sc_s8je 0 (sc2_s8jh `cast` ...)
                  of s'#_a5GP { __DEFAULT ->
                  $s$wa_s8ji (+# sc_s8je 1) (+# sc1_s8jf 1) (s'#_a5GP `cast` ...)
                  }
              }; } in
        case $s$wa_s8ji 0 0 (ipv_a5g3 `cast` ...)
        of _ { (# ipv6_a4MX, ipv7_a4MY #) ->
        case ipv7_a4MY of _ { I# dt4_a5xy ->
        -- end of vector creation

        letrec {
          a2_s7Q6 :: Int# -> State# RealWorld -> (# State# RealWorld, () #)
          a2_s7Q6 =
            \ (x_a5HT :: Int#) (eta_B1 :: State# RealWorld) ->
              letrec {
                a3_s7Q5 :: Int# -> State# RealWorld -> (# State# RealWorld, () #)
                a3_s7Q5 =
                  \ (x1_X5J9 :: Int#) (eta1_XP :: State# RealWorld) ->
                    letrec {
                      a4_s7MZ :: Int# -> State# RealWorld -> (# State# RealWorld, () #)
                      a4_s7MZ =
                        \ (x2_X5Jl :: Int#) (s1_X4Xb :: State# RealWorld) ->
                          case writeIntArray#
                                 (ipv1_a5g4 `cast` ...)
                                 (+# (*# x1_X5J9 100) x2_X5Jl)
                                 1
                                 (s1_X4Xb `cast` ...)
                          of s'#_a5GP { __DEFAULT ->

                          -- the interesting part! ------------------
                          case x2_X5Jl of wild_X1y {
                            __DEFAULT -> a4_s7MZ (+# wild_X1y 1) (s'#_a5GP `cast` ...);
                            99 ->
                              case x1_X5J9 of wild1_X1o {
                                __DEFAULT -> a3_s7Q5 (+# wild1_X1o 1) (s'#_a5GP `cast` ...);
                                99 ->
                                  case x_a5HT of wild2_X1c {
                                    __DEFAULT -> a2_s7Q6 (+# wild2_X1c 1) (s'#_a5GP `cast` ...);
                                    99999 -> (# s'#_a5GP `cast` ..., () #)
                                  }
                              }
                          }
                          }; } in
                    a4_s7MZ 0 eta1_XP; } in
              a3_s7Q5 0 eta_B1; } in
        a2_s7Q6 0 (ipv6_a4MX `cast` ...)
        }
        }
        };
      True ->
        case error
               (unpackAppendCString#
                  "Primitive.basicUnsafeNew: length to large: "#
                  (case $wshowSignedInt 0 10000 ([])
                   of _ { (# ww5_a5wO, ww6_a5wP #) ->
                   : ww5_a5wO ww6_a5wP
                   }))
        of wild_00 {
        }
    }
    }

main :: IO ()
main = main1 `cast` ...

main2 :: State# RealWorld -> (# State# RealWorld, () #)
main2 = runMainIO1 (main1 `cast` ...)

main :: IO ()
main = main2 `cast` ...

我不得不承认，我基本上不知道为什么有些代码会避免创建堆栈帧，而有些则不会。我怀疑从“内部”向外内联会有所帮助，并且快速检查告诉我Control.Monad.Loop 使用 CPS 编码，这可能与此处相关，尽管Monad.Loop 解决方案对浮动很敏感，而我不能立即从核心确定为什么 test_c with let floating 无法在单个堆栈帧中运行。

现在，在单个堆栈帧中运行的性能优势很小。我们已经看到test_b 只比test_a 快一点。我在答案中加入了这条弯路，因为我发现它很有启发性。

状态破解和严格绑定

所谓的state hack 使得 GHC 积极地内联到 IO 和 ST 动作。我想我应该在这里提一下，因为除了让浮动之外，这是另一个可以彻底破坏性能的东西。

通过优化-O 启用状态破解，并且可能会渐近地减慢程序速度。来自Reid Barton的一个简单例子：

import Control.Monad
import Debug.Trace

expensive :: String -> String
expensive x = trace "$$$" x

main :: IO ()
main = do
  str <- fmap expensive getLine
  replicateM_ 3 $ print str

在 GHC-7.10.2 中，这会打印一次 "$$$" 而不进行优化，但会打印三次 -O2。而且似乎在 GHC-7.10 中，我们无法通过 -fno-state-hack（这是 Reid Barton 的链接票证的主题）摆脱这种行为。

严格的一元绑定可靠地解决了这个问题：

main :: IO ()
main = do
  !str <- fmap expensive getLine
  replicateM_ 3 $ print str

我认为在 IO 和 ST 中进行严格绑定是一个好习惯。而且我有一些经验（虽然不是确定的；我远不是 GHC 专家）如果我们使用 -fno-full-laziness，则特别需要严格的绑定。显然，完全的懒惰可以帮助摆脱由 state hack 引起的内联引入的一些工作重复；使用test_b 并且没有完全的懒惰，省略对!mvec <- V.unsafeThaw vec 的严格绑定会导致轻微的减速和极其丑陋的核心输出。

Answer 2

根据我的经验，forM_ [0..n-1] 可以表现良好，但不幸的是它并不可靠。只需将 INLINE 杂注添加到 test_a 并使用 -O2 就可以让它运行得更快（对我来说是 4 秒到 1 秒），但是手动内联它（复制粘贴）会再次减慢速度。

一个更可靠的函数是for from statistics 实现为

-- | Simple for loop.  Counts from /start/ to /end/-1.
for :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
for n0 !n f = loop n0
  where
    loop i | i == n    = return ()
           | otherwise = f i >> loop (i+1)
{-# INLINE for #-}

使用它看起来类似于带有列表的forM_：

test_d :: MV.IOVector Int -> IO ()
test_d mv =
  for 0 times $ \_ ->
    for 0 side $ \i ->
      for 0 side $ \j ->
        MV.unsafeWrite mv (i*side + j) 1

但性能可靠（对我来说是 0.85 秒），没有任何分配列表的风险。