从 ByteString 创建 Repa

Repa creation from ByteString

最初我有一个 ByteString,然后将其解压缩并转换为 Int16,这部分过程花费的时间相对较少。然后,我使用以下行将 Int16 列表转换为 Repa 数组,

Repa.fromListUnboxed (Z :. bytesOfDataPerImage `div` 2) listOfInts

根据探查器,此行占用了 CPU 的 ~40% 时间,这可能只是表明我正在执行的计算不保证使用 Repa。从 ByteString 到 Repa 数组时是否有更有效的途径?

我试过 Repa fromByteString 函数,虽然转换

Array B DIM1 Word8 -> Array U DIM1 Int16

非常慢。我通过首先将数组重塑为 Word8 的二维数组,然后折叠成 Int16 来执行此操作。也许字节数组是正确的方法,而我的转换方法是错误的。

convertImageData :: Array B DIM1 Word8 -> Array U DIM1 Int16
convertImageData !arr = Repa.foldS convertWords 0 (Repa.map fromIntegral (splitArray arr))

splitArray :: Array B DIM1 Word8 -> Array U DIM2 Word8
splitArray !arr = computeUnboxedS $ reshape (Z :. ((size $ extent arr) `div` 2) :. 2) arr


convertWords :: Int16 -> Int16 -> Int16
convertWords !word1 !word2 = (word1 `shiftL` 8) .|. word2

在某些情况下,该程序正在针对用 C/C++ 编写的同一程序进行基准测试。

您最初转换为列表然后调用 Repa.fromListUnboxed 的方法肯定非常慢,因为您所做的只是强制列表中的元素,而不是按顺序将其加载到未装箱的数组中。这就是为什么转换成列表只需要很少的时间,因为它所做的只是创建一堆 thunk,但实际的计算是在将它加载到数组中时发生的。

你的第二种方法肯定更好,但仍然有不必要的步骤,例如。不需要 reshape 数组,您只需将新大小传递给 fromByteString 函数即可。所以这是一个稍微改进的版本:

bytesToRepaOriginal :: Bytes.ByteString -> Repa.Array Repa.U Repa.DIM1 Int16
bytesToRepaOriginal bs =
  Repa.foldS
    convertWords
    0
    (Repa.map fromIntegral $
     Repa.fromByteString (Repa.Z Repa.:. (Bytes.length bs `div` 2) Repa.:. 2) bs)
Repa 中的

fromByteString 函数和 B 表示由于某种原因并不是特别快,所以有一种更快的方法,即通过直接索引 [= 来构造一个数组19=]:

bytesToRepaP :: Monad m => Bytes.ByteString -> m (Repa.Array Repa.U Repa.DIM1 Int16)
bytesToRepaP bs =
  Repa.computeUnboxedP $
  Repa.fromFunction
    (Repa.Z Repa.:. (Bytes.length bs `div` 2))
    (\(Repa.Z Repa.:. i) ->
       let i' = i * 2
           f = Bytes.unsafeIndex bs
        in (fromIntegral (f i') `shiftL` 8) .|. fromIntegral (f (i' + 1)))

使用 Repa.computeUnboxedS 切换到顺序计算会使您的速度降低 x2 倍,但由于我们正在尝试对其进行优化,因此我们需要一直进行并行计算。

不是要从一个非常好的 Repa 库中窃取所有风头,我还想展示所有这些如何与新的 massiv 库一起工作:

import Data.Massiv.Array as Massiv

bytesToMassiv :: Bytes.ByteString -> Massiv.Array Massiv.U Massiv.Ix1 Int16
bytesToMassiv bs =
  Massiv.makeArrayR U Par (Bytes.length bs `div` 2)
    (\i ->
       let i' = i * 2
           f = Bytes.unsafeIndex bs
        in (fromIntegral (f i') `shiftL` 8) .|. fromIntegral (f (i' + 1)))

只是为了展示一些具体的数字来展示优化的实际效果,这是一个精简的标准基准:

benchmarking fromByteString/Massiv Parallel
time                 1.114 ms   (1.079 ms .. 1.156 ms)

benchmarking fromByteString/Repa Parallel
time                 1.954 ms   (1.871 ms .. 2.040 ms)

benchmarking fromByteString/Repa Original
time                 15.80 ms   (15.67 ms .. 15.92 ms)