未装箱向量中没有与 unsafeUpdate_ 的流融合答案

【问题标题】：No stream fusion with unsafeUpdate_ in unboxed vector未装箱向量中没有与 unsafeUpdate_ 的流融合
【发布时间】：2013-06-08 03:51:42
【问题描述】：

如果使用unsafeUpdate_ 函数更新vector 的某些元素，是否可以在处理vector 时保持流融合？在我做的测试中，答案似乎是否定的。对于下面的代码，在upd函数中生成了临时向量，在核心中确认：

module Main where
import Data.Vector.Unboxed as U

upd :: Vector Int -> Vector Int
upd v = U.unsafeUpdate_ v (U.fromList [0]) (U.fromList [2])

sum :: Vector Int -> Int
sum = U.sum . upd

main = print $ Main.sum $ U.fromList [1..3]

在核心中，$wupd 函数用于sum - 如下所示，它生成新的bytearray：

$wupd :: Vector Int -> Vector Int
$wupd =
  \ (w :: Vector Int) ->
    case w `cast` ... of _ { Vector ipv ipv1 ipv2 ->
    case main11 `cast` ... of _ { Vector ipv3 ipv4 ipv5 ->
    case main7 `cast` ... of _ { Vector ipv6 ipv7 ipv8 ->
    runSTRep
      (\ (@ s) (s :: State# s) ->
         case >=# ipv1 0 of _ {
           False -> case main6 ipv1 of wild { };
           True ->
             case newByteArray# (*# ipv1 8) (s `cast` ...)
             of _ { (# ipv9, ipv10 #) ->
             case (copyByteArray# ipv2 (*# ipv 8) ipv10 0 (*# ipv1 8) ipv9)
                  `cast` ...

sum 函数的核心中有一个很好的紧密循环，但就在该循环之前，有一个对 $wupd 函数的调用，因此是一个临时生成。

这里的例子有没有办法避免临时生成？我的想法是，更新索引 i 中的向量是解析流但仅作用于索引 i 中的流（跳过其余部分），并将那里的元素替换为另一个元素的情况。所以，在任意位置更新向量不应该破坏流融合，对吧？

【问题讨论】：

旁白：yarr 支持融合和展开以及就地更改：hackage.haskell.org/packages/archive/yarr/1.3.1/doc/html/…
@leventov，非常好！由于您似乎是该库的作者，您能否分享如何为上面的矢量示例编写就地可熔更新函数？当然，你为 yarr 所做的应该适用于这里，不是吗？
大，yarr-1.* 的大缺点它与 IO 紧密相连。就地编辑只是do { ... write arr 0 2; ... }。

标签： haskell vector stream fusion

【解决方案1】：

我不能 100% 确定，因为 vector 一直是乌龟（你永远不会真正到达实际的实现，总是有另一个间接），但据我了解，update变体通过克隆强制一个新的临时：

unsafeUpdate_ :: (Vector v a, Vector v Int) => v a -> v Int -> v a -> v a
{-# INLINE unsafeUpdate_ #-}
unsafeUpdate_ v is w
  = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))

unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
{-# INLINE unsafeUpdate_stream #-}
unsafeUpdate_stream = modifyWithStream M.unsafeUpdate

和modifyWithStream 调用clone（和new），

modifyWithStream :: Vector v a
                 => (forall s. Mutable v s a -> Stream b -> ST s ())
                 -> v a -> Stream b -> v a
{-# INLINE modifyWithStream #-}
modifyWithStream p v s = new (New.modifyWithStream p (clone v) s)

new :: Vector v a => New v a -> v a
{-# INLINE_STREAM new #-}
new m = m `seq` runST (unsafeFreeze =<< New.run m)

-- | Convert a vector to an initialiser which, when run, produces a copy of
-- the vector.
clone :: Vector v a => v a -> New v a
{-# INLINE_STREAM clone #-}
clone v = v `seq` New.create (
  do
    mv <- M.new (length v)
    unsafeCopy mv v
    return mv)

我认为vector 不可能再次摆脱unsafeCopy。

【讨论】：

根据vector recycling paper，clone/new 应该融合在一起，尽管我在实验中看不到它。

【解决方案2】：

如果您需要更改一个或很少的元素，repa 和 yarr 库中有很好的解决方案。它们保留了融合（我不确定 repa）和 Haskell 惯用语。

Repa，使用fromFunction:

upd arr = fromFunction (extent arr) ix
  where ix (Z .: 0) = 2
        ix i = index arr i

Yarr，使用Delayed：

upd arr = Delayed (extent arr) (touchArray arr) (force arr) ix
  where ix 0 = return 2
        ix i = index arr i

【讨论】：