Haskell
vector

動的計画法にData.Vector.constructNは使うべきではない。


constructN

haskellのvectorパッケージにはconstructNという関数がある。

-- | /O(n)/ Construct a vector with @n@ elements by repeatedly applying the

-- generator function to the already constructed part of the vector.
--
-- > constructN 3 f = let a = f <> ; b = f <a> ; c = f <a,b> in f <a,b,c>
--
constructN :: Int -> (Vector a -> a) -> Vector a

ソース

コメントからわかるように、constructNは引数として整数nと配列の要素を生成する関数fを受け取って長さnの配列を生成する。

この時、生成関数fは「0からn-1番目まで生成された配列」から「n番目の要素」を返すようにする。

さてこの関数を使うと動的計画法が簡単に書ける。例えばフィボナッチ数の入った配列なら、

fib :: Int -> Vector Integer

fib n = constructN n gen
where
gen vec | i <= 1 = i
| otherwise = vec ! (i - 1) + vec !
(i - 2)
where i = length vec


サンクの呪い

しかしながらこのfibは極めて遅い。

原因を調べてみると、Data.VectorconstructNは要素を遅延評価することがわかった。

すなわち、要素がa,b,cであるvectorをV<a,b,c>と書くことにすると、

constructN n gen = V<v_0,...,v_{n-1}>

where
v_0 = gen V<>
v_1 = gen V<v_0>
v_2 = gen V<v_0, v_1>
...
v_{n-1} = gen V<v_0, v_1, ..., v_{n-2}

のようなvectorが作られる。軽く地獄絵図である。

このvectorがどの程度メモリを消費するか考えてみよう。

Vector aは配列のスライスである。

data Vector a = Vector {-# UNPACK #-} !Int

{-# UNPACK #-} !Int
{-# UNPACK #-} !(Array a)

ソース

Array a型はunpackしても1wordのようなので、Vector aのサイズは3 wordsとなる。これにサンクのメモリオーバーヘッド(2 words)を加えると、5 words = 40 bytesも消費してしまうのだ。aIntの場合、正格に評価すれば2 wordsで済むので2.5倍のメモリオーバーヘッドが生じる。

実際に次のようなプログラムで確かめてみよう。


Main.hs

{-# LANGUAGE BangPatterns #-}

import Data.Vector
import qualified Data.Vector.Unboxed as U
import Prelude hiding(length)
import Control.Exception
import System.Environment
import Control.Monad

{-# NOINLINE fib #-}
fib :: Int -> Vector Int
fib n = constructN n gen
where gen v | i <= 1 = i
| otherwise = v ! (i - 1) + v ! (i - 2)
where i = length v

{-# NOINLINE fibU #-}
fibU :: Int -> U.Vector Int
fibU n = U.constructN n gen
where gen v | i <= 1 = i
| otherwise = v U.! (i - 1) + v U.! (i - 2)
where i = U.length v

{-# NOINLINE fibS #-}
fibS :: Int -> Vector Int
fibS n = unfoldrN n gen (0, 1)
where gen (f1, f2) = Just (f1, (f2,(f1 + f2)))

main :: IO ()
main = do
args <- getArgs
case args of
["fib"] -> void $ evaluate (fib 1000000)
["fibS"] -> void $ evaluate (fibS 1000000)
["fibU"] -> void $ evaluate (fibU 1000000)
_ -> pure ()


fibconstructNを用いたバージョン、fibSunfoldrNを用いて正格にしたバージョン、fibUData.Vector.Unboxを使用したバージョンである。

それぞれ長さ1,000,000のvectorを作る。

fibSが中間データのタプルを生成しないようにするため、ghc -O2 -rtsopts Main.hsでコンパイルする。

$ ./Main +RTS -s -RTS fib

48,050,576 bytes allocated in the heap
85,079,672 bytes copied during GC
40,899,272 bytes maximum residency (4 sample(s))
525,624 bytes maximum slop
72 MB total memory in use (0 MB lost due to fragmentation)
...
$ ./Main +RTS -s -RTS fibS
24,050,984 bytes allocated in the heap
41,905,160 bytes copied during GC
19,531,968 bytes maximum residency (3 sample(s))
370,496 bytes maximum slop
31 MB total memory in use (0 MB lost due to fragmentation)
...
$ ./Main +RTS -s -RTS fib
8,043,448 bytes allocated in the heap
7,128 bytes copied during GC
36,064 bytes maximum residency (1 sample(s))
376,672 bytes maximum slop
9 MB total memory in use (0 MB lost due to fragmentation)
...

このようにfibは48MB, fibSは24MB, fibUは8MBを消費している。

- fibUについては、1,000,000個のInt#型の配列のサイズがちょうど8MB消費するので、最小限のメモリしか使用しないことがわかる。

- fibSについても、1,000,000個のInt型の値のサイズが16MB、それら1,000,000個のポインタを格納する配列のサイズが8MBであるので、これも最小限である。

- 一方でfibはサンクに40MBを消費していることがわかる。この値は先ほど見積もった(5 words * 8 * 1,000,000)とほぼ一致する。


原因

なぜData.Vector.constructNは遅延評価されるのか見ていこう。

constructN :: Int -> (Vector a -> a) -> Vector a

{-# INLINE constructN #-}
constructN = G.constructN

ソース

なるほど、Data.Vector.Generic.constructNが実装のようだ。

constructN :: forall v a. Vector v a => Int -> (v a -> a) -> v a

{-# INLINE constructN #-}
-- NOTE: We *CANNOT* wrap this in New and then fuse because the elements
-- might contain references to the immutable vector!
constructN !n f = runST (
do
v <- M.new n
v' <- unsafeFreeze v
fill v' 0
)
where
fill :: forall s. v a -> Int -> ST s (v a)
fill !v i | i < n = let x = f (unsafeTake i v)
in
elemseq v x $
do
v' <- unsafeThaw v
M.unsafeWrite v' i x
v'' <- unsafeFreeze v'
fill v'' (i+1)

fill v _ = return v

ソース

なるほど、fill関数内で生成関数fを使って配列を初期化している。ここでelemseq v xというのが気になるところだ。

class MVector (Mutable v) a => Vector v a where

(中略)
-- | Evaluate @a@ as far as storing it in a vector would and yield @b@.
-- The @v a@ argument only fixes the type and is not touched. The method is
-- only used for optimisation purposes. Thus, it is safe for instances of
-- 'Vector' to evaluate @a@ less than it would be when stored in a vector
-- although this might result in suboptimal code.
--
-- > elemseq v x y = (singleton x `asTypeOf` v) `seq` y
--
-- Default defintion: @a@ is not evaluated at all
--
elemseq :: v a -> a -> b -> b

{-# INLINE elemseq #-}
elemseq _ = \_ x -> x

ソース

なるほど、デフォルトでは何もしてくれないらしい。

ではインスタンス宣言では何か定義してくれているだろうか。

instance G.Vector Vector a where

{-# INLINE basicUnsafeFreeze #-}
basicUnsafeFreeze (MVector i n marr)
= Vector i n `liftM` unsafeFreezeArray marr

{-# INLINE basicUnsafeThaw #-}
basicUnsafeThaw (Vector i n arr)
= MVector i n `liftM` unsafeThawArray arr

{-# INLINE basicLength #-}
basicLength (Vector _ n _) = n

{-# INLINE basicUnsafeSlice #-}
basicUnsafeSlice j n (Vector i _ arr) = Vector (i+j) n arr

{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM (Vector i _ arr) j = indexArrayM arr (i+j)

{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy (MVector i n dst) (Vector j _ src)
= copyArray dst i src j n

ソース

そんなことはなかった

一方、Data.Vector.Unbox等ではelemseq _ = seqと定義されており、elemseq v x bxを正格に評価してくれる。


回避策

elemseq _ = seqとなるG.Vector v aのインスタンスを作ってやれば良い。多少強引になるが以下のように定義できる。

newtype Strict a = Strict { fromStrict :: a } deriving(Eq,Ord)

instance {-# Overlaps #-} G.Vector Vector (Strict a) where
{-# INLINE basicUnsafeFreeze #-}
basicUnsafeFreeze :: forall m. PrimMonad m => MV.MVector (PrimState m) (Strict a) -> m (V.Vector (Strict a))
basicUnsafeFreeze v = unsafeCoerce (G.basicUnsafeFreeze (unsafeCoerce v :: MV.MVector (PrimState m) a) :: m (V.Vector a))
{-# INLINE basicUnsafeThaw #-}
basicUnsafeThaw :: forall m. PrimMonad m => V.Vector (Strict a) -> m (MV.MVector (PrimState m) (Strict a))
basicUnsafeThaw v = unsafeCoerce $ (G.basicUnsafeThaw (unsafeCoerce v :: V.Vector a) :: m (MV.MVector (PrimState m) a))
{-# INLINE basicLength #-}
basicLength v = G.basicLength (unsafeCoerce v :: V.Vector a)
{-# INLINE basicUnsafeSlice #-}
basicUnsafeSlice offset len v = unsafeCoerce $ G.basicUnsafeSlice offset len (unsafeCoerce v :: V.Vector a)
{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM :: forall m. Monad m => V.Vector (Strict a) -> Int -> m (Strict a)
basicUnsafeIndexM v i = unsafeCoerce $ (G.basicUnsafeIndexM (unsafeCoerce v :: V.Vector a) i :: m a)
{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy :: forall m. PrimMonad m => MV.MVector (PrimState m) (Strict a) -> V.Vector (Strict a) -> m ()
basicUnsafeCopy dst src = G.basicUnsafeCopy (unsafeCoerce dst :: MV.MVector (PrimState m) a) (unsafeCoerce src :: V.Vector a)
{-# INLINE elemseq #-}
elemseq _ = seq

Strict Int型を使えば、Data.Vector.Generic.constructNも正格になる。(Data.Vector.constructNは元々のinstance定義を使ってしまうので正格にならないことに注意)

{-# NOINLINE fibG #-}

fibG :: Int -> Vector Int
fibG n = unsafeCoerce $ (G.constructN n (Strict . gen) :: Vector (Strict Int))
where gen v | i <= 1 = i
| otherwise = fromStrict (v ! (i - 1)) + fromStrict (v ! (i - 2))
where i = length v

このfibGだと長さ1,000,000の配列生成でメモリ消費が24MBとなる。

$ ./Main +RTS -s -RTS fibG

24,050,952 bytes allocated in the heap
29,207,520 bytes copied during GC
20,531,440 bytes maximum residency (3 sample(s))
374,544 bytes maximum slop
33 MB total memory in use (0 MB lost due to fragmentation)


結論

Data.Vector.Vectorは非正格なので取り扱いに注意すべし。

Data.Vector.Strictみたいなものがあれば良いのだが。