atomicModifyIORef' の歴史 - あどけない話

遅延評価、何も分からん。
— 山本和彦 (@kazu_yamamoto) 2023年2月18日

Haskellの並行プログラミングの奥義であるatomicModifyIORef'の歴史。内容は予告なく修正される。

atomicModifyIORef'は「IORef」と「関数」を引数に取る。その関数は、IORefが現在指している「古い値」を受け取り、IORefが新たに指すべき「新しい値」と「返り値」をタプルで返すように書く。

atomicModifyIORef' :: IORef a -> (a -> (a, b)) -> IO b

新しい値は、サンクが作られた時点で、古い値と不可分にスワップされ、その後評価される。評価は重くとも、サンクを作ることは軽いので、スワップは多くの場合成功する。

`atomicModifyIORef` の問題点

もうすっかり忘れていたが、Mightyの開発過程で、atomicModifyIORefがスペースリークをすることを発見し、解決策を2011年の Monad.Reader issue 19に執筆した。

以下は、時間の文字列を1秒に一回作ってキャッシュしておくコード：

atomicModifyIORef ref (\_ -> (tmstr, ()))

サンクが潰れずにスペースリークとなる。解決策として、以下を考案。

x <- atomicModifyIORef ref (\_ -> (tmstr, ()))
x ‘seq‘ return ()

これを Monad.Reader に書いた。その後、バンパターンを使うともっとカッコよく書けることが判明。

!_ <- atomicModifyIORef ref (\_ -> (tmstr, ()))

注意深くコードを読んだ人は疑問に思うだろう。タプルの右側の要素しか評価していないのに、どうして左側の新しい値に起因するスペースリークがなくなるか？

このコードでは、新しく作られる値のサンクを潰すことはできない。しかし、右の要素を評価することでタプルが生成される。すると、左側のサンクは、古い値を参照していないことが分かり、GCが回収できる。

このコードは、IORefが指す古い値と返り値を使わないので、現在では atomicWriteIORef が利用できる。

`atomicModifyIORef'` の登場

Monad.Reader issue 19をきっかけに、Joey Adams氏が atomicModifyIORef' を提案。以下は、2011年にマージされたAdd strict versions of modifyIORef and atomicModifyIORefより：

atomicModifyIORef' :: IORef a -> (a -> (a,b)) -> IO b
atomicModifyIORef' ref f = do
    b <- atomicModifyIORef ref
            (\x -> let (a, b) = f x
                    in (a, a `seq` b))
    b `seq` return b

baseライブラリ4.6.0以降で atomicModifyIORef'が利用可能となった。つまり、2012年にリリースされた GHC 7.6.1 から利用できる。

Simon Marlow氏が、2013年に発行されたParallel and Concurrent Programming in Haskell の284ページでこのパターンを紹介。atomicModifyIORef' 自体への言及はなし。2014年に発行されたHaskellによる並列・並行プログラミングの363ページでは、僕が atomicModifyIORef' に関する訳注を入れた。

参照はサンクが作られた時点でスワップされて、そのサンクは atomicModifyIORef'の終了後、以下のように評価される。

atomicModifyIORefの外側で b を返す前に、bをWHNFにしようとする
bを評価するには、内部で仕込まれたseqにより、aもWHNFになる
その後bがWHNFになる。

このころはプリミティブな命令として atomicModifyMutVar# があり、atomicModifyIORef は、そのラッパーだった。以下は不可分性はないが、atomicModifyIORef の動作をエミュレートした疑似コードである。

pseudoAtomicModifyIORef :: IORef a -> (a -> (a, b)) -> IO b
pseudoAtomicModifyIORef ref f = do
  x <- readIORef ref
  case f x of
    fx -> do  -- thunk 'fx'
      -- only CAS success case
      writeIORef ref $ fst fx -- thunk 'fst fx'
      putStrLn "swap"
      return $ snd fx -- thunk 'snd fx'

これを使い Joey Adams氏の実装をエミュレートする：

joey :: IORef a -> (a -> (a, b)) -> IO b
joey ref f = do
    b <- pseudoAtomicModifyIORef ref
            (\x -> let (a, b) = f x
                    in (a, a `seq` b))
    b `seq` return b

また、評価順序を知るために、たくさん trace 仕込んだ f を用意する：

f :: Int -> (Int, ())
f n = (trace "tupple" ((trace "new" n + 1), (trace "ret" ())))

GHCiで試してみよう。

> ref <-newIORef (0::Int)
> joey ref f
swap
tupple
new
ret

"new" よりも前に "swap" が表示されているのが分かる。

atomicModifyIORef' にどんなに正格な関数を与えても、サンクが先に IORef に格納されて、その後で新しい値が計算される。以下のような正格な関数を考える：

f' :: Int -> (Int, ())
f' n = let !n' = trace "new" n + 1 in (trace "tupple" (n', (trace "ret" ())))

使ってみよう。

> joey ref f'
swap
new
tupple
ret

"tupple" よりも "new" の方が早く生成されるが、それでも "swap" の後だと分かる。

性能向上

オリジナルの atomicModifyIORef' は、タプルでパターンマッチした後、新しいタプルを作成している。これは無駄である。parcs 氏は、セレクターサンクを使って性能を向上させること提案した。

atomicModifyIORef' ref f = do
    b <- atomicModifyIORef ref $ \a ->
            case f a of
                v@(a',_) -> a' `seq` v
    b `seq` return b

このコードはタプルを生成しない。case文はフィールドを取り出すために使われている。これは、セレクターサンクとなって、GCが高速に扱えるらしい。

しかし、本当にこの方法で、新しい値が評価される前にスワップされるのだろうか？ v を返す前に a' が評価されているように見える！

違う。違う。あなたの解釈は違うのだ。Joey Adams氏の実装は、タプルの遅延性を利用しているからうまくいくと思っているなら、違うのだ。atomicModifyMutVar# が遅延させるのは、関数(ラムダ式)全体であって、タプルの生成には頼っていない。

先ほど、正格な関数 f' を与えても、うまく動作していたことを思い出そう。

parcs氏の実装をエミュレートする：

parcs ref f = do
    b <- pseudoAtomicModifyIORef ref $ \a ->
            case f a of
                v@(a',_) -> a' `seq` v
    b `seq` return b

使ってみる。

> ref <-newIORef (0::Int)
> parcs ref f'
swap
new
tupple
ret

やはり、"swap" が最初に起きていることが分かる。

`atomicModifyMutVar2#`

David Feuer氏が、ghc proposals 0149を出し、atomicModifyIORefのスペースリークをなくすために、新たなプリミティブである atomicModifyMutVar2# の作成を提案した。

atomicModifyMutVar では、第一フィールドと第二フィールド両方にサンクを作成していた。atomicModifyMutVar2#では、第一フィールドだけにサンクを作る。第二フィールドは正格に評価されるので「サンクの割り当て直後に潰す」という無駄が生じない。

atomicModifyMutVar2# :: MutVar# s a -> (a -> c) -> State# s -> (# State# s, a, c #)

引数の関数の型から、タプルを仮定してないことが分かる。これを一般的なHaskellのコードから呼ぶための関数にするラッパーが、2つある。

atomicModifyIORef2Lazy :: IORef a -> (a -> (a,b)) -> IO (a, (a, b))
atomicModifyIORef2Lazy (IORef (STRef r#)) f =
  IO (\s -> case atomicModifyMutVar2# r# f s of
              (# s', old, res #) -> (# s', (old, res) #))

atomicModifyIORef2 :: IORef a -> (a -> (a,b)) -> IO (a, (a, b))
atomicModifyIORef2 ref f = do
  r@(_old, (_new, _res)) <- atomicModifyIORef2Lazy ref f
  return r

なぜ2つに分かれているのかは不明。現在のatomicModifyIORef'の実装はこう：

atomicModifyIORef' :: IORef a -> (a -> (a,b)) -> IO b
atomicModifyIORef' ref f = do
  (_old, (_new, !res)) <- atomicModifyIORef2 ref $
    \old -> case f old of
       r@(!_new, _res) -> r
  pure res

以下は、日比野さんが書いた検証コード：

pseudoModifyIORef2Lazy :: IORef a -> (a -> (a, b)) -> IO (a, (a, b))
pseudoModifyIORef2Lazy ref f = do
  x <- readIORef ref
  case f x of
    fx -> do
      -- only CAS success case
      writeIORef ref (fst fx)
      putStrLn "swap"
      return (x, fx)

pseudoModifyIORef2 :: IORef a -> (a -> (a,b)) -> IO (a, (a, b))
pseudoModifyIORef2 ref f = do
  r@(_old, (_new, _res)) <- pseudoModifyIORef2Lazy ref f
  return r

pseudoModifyIORef :: IORef a -> (a -> (a,b)) -> IO b
pseudoModifyIORef ref f = do
  (_old, ~(_new, res)) <- pseudoModifyIORef2 ref f
  pure res

david :: IORef a -> (a -> (a,b)) -> IO b
david ref f = do
    (_old, (_new, !res)) <- pseudoModifyIORef2 ref $
        \old -> case f old of
            r@(!_new, _res) -> r
    pure res

動作順序を調べてみる：

> ref <-newIORef (0::Int)
> david ref f'
swap
new
tupple
ret

"swap" が一番最初に出力されているのが分かる。

atomicModifyIORef の問題点

atomicModifyIORef' の登場

性能向上

atomicModifyMutVar2#

`atomicModifyIORef` の問題点

`atomicModifyIORef'` の登場

`atomicModifyMutVar2#`