编写一次并行数组 Haskell 表达式，在 CPU 和 GPU 上运行 repa 并加速

Posted 2023-03-06

【中文标题】编写一次并行数组 Haskell 表达式，在 CPU 和 GPU 上运行 repa 并加速

【英文标题】：Write a parallel array Haskell expression once, run on CPUs & GPUs with repa and accelerate

【发布时间】：2014-04-21 17:13:25

【问题描述】：

修复和加速 API 相似性

Haskell repa 库用于在 CPU 上自动进行并行数组计算。加速库是 GPU 上的自动数据并行性。 API 非常相似，具有相同的 N 维数组表示。甚至可以使用fromRepa 和toRepa 在加速和repa 数组之间切换Data.Array.Accelerate.IO：

fromRepa :: (Shapes sh sh', Elt e) => Array A sh e -> Array sh' e
toRepa   :: Shapes sh sh'          => Array sh' e  -> Array A sh e

有多个用于加速的后端，包括 LLVM、CUDA 和 FPGA（参见 http://www.cse.unsw.edu.au/~keller/Papers/acc-cuda.pdf 的图 2）。我发现了一个用于加速的repa backend，尽管该库似乎没有得到维护。鉴于 repa 和加速编程模型相似，我希望有一种优雅的方式在它们之间切换，即编写一次的函数可以使用 repa 的 R.computeP 或加速的后端之一执行，例如使用 CUDA run 函数。

两个非常相似的函数：南瓜上的 Repa 和 Accelerate

取一个简单的图像处理阈值函数。如果灰度像素值小于 50，则将其设置为 0，否则保留其值。这是它对南瓜的作用：

以下代码展示了repa和加速实现：

module Main where

import qualified Data.Array.Repa as R
import qualified Data.Array.Repa.IO.BMP as R
import qualified Data.Array.Accelerate as A
import qualified Data.Array.Accelerate.IO as A
import qualified Data.Array.Accelerate.Interpreter as A

import Data.Word

-- Apply threshold over image using accelerate (interpreter)
thresholdAccelerate :: IO ()
thresholdAccelerate = do
  img <- either (error . show) id `fmap` A.readImageFromBMP "pumpkin-in.bmp"
  let newImg = A.run $ A.map evalPixel (A.use img)
  A.writeImageToBMP "pumpkin-out.bmp" newImg
    where
      -- *** Exception: Prelude.Ord.compare applied to EDSL types
      evalPixel :: A.Exp A.Word32 -> A.Exp A.Word32
      evalPixel p = if p > 50 then p else 0

-- Apply threshold over image using repa
thresholdRepa :: IO ()
thresholdRepa = do
  let arr :: IO (R.Array R.U R.DIM2 (Word8,Word8,Word8))
      arr = either (error . show) id `fmap` R.readImageFromBMP "pumpkin-in.bmp" 
  img <- arr
  newImg <- R.computeP (R.map applyAtPoint img)
  R.writeImageToBMP "pumpkin-out.bmp" newImg
  where
    applyAtPoint :: (Word8,Word8,Word8) -> (Word8,Word8,Word8)
    applyAtPoint (r,g,b) =
        let [r',g',b'] = map applyThresholdOnPixel [r,g,b]
        in (r',g',b')
    applyThresholdOnPixel x = if x > 50 then x else 0

data BackendChoice = Repa | Accelerate

main :: IO ()
main = do
  let userChoice = Repa -- pretend this command line flag
  case userChoice of
    Repa       -> thresholdRepa
    Accelerate -> thresholdAccelerate

问题：我可以只写一次吗？

thresholdAccelerate 和thresholdRepa 的实现非常相似。是否有一种优雅的方法可以编写一次数组处理函数，然后以编程方式在交换机中选择多核 CPU（repa）或 GPU（加速）？我可以考虑根据我想要 CPU 还是 GPU 来选择我的导入，即导入 Data.Array.Accelerate.CUDA 或 Data.Array.Repa 以执行 Acc a 类型的操作：

run :: Arrays a => Acc a -> a

或者，使用类型类，例如大致是这样的：

main :: IO ()
main = do
  let userChoice = Repa -- pretend this is a command line flag
  action <- case userChoice of
    Repa       -> applyThreshold :: RepaBackend ()
    Accelerate -> applyThreshold :: CudaBackend ()
  action

或者，对于我希望为 CPU 和 GPU 表达的每个并行数组函数，我必须实现它两次 --- 一次使用 repa 库，另一次使用加速库？

【问题讨论】：

Accelerate 有一个 repa 后端：github.com/blambo/accelerate-repa - 从未使用过，但我很高兴知道这是否足够有趣。

对，似乎打算使用 repa 后端并使用 accelerate 编写代码，但我想无论出于何种原因，对此都没有太大兴趣。

是的，但我认为这可能是合适的，因为accelerate 允许我们编写一次函数并让它根据您的需要在并行 CPU 或 GPU 上运行。这是一个模块切换，而不是翻转构造函数，但您可以在前者之上编写后者。

【参考方案1】：

简短的回答是，目前很遗憾，您需要编写两个版本。

不过，我们正在努力为 Accelerate 提供 CPU 支持，这将消除对代码的 Repa 版本的需求。特别是，Accelerate 最近获得了一个新的基于 LLVM 的后端，它同时针对 GPU 和 CPU：https://github.com/AccelerateHS/accelerate-llvm

这个新的后端仍然不完整、有缺陷且处于试验阶段，但我们正计划使其成为当前 CUDA 后端的可行替代方案。

【讨论】：

太棒了！ w.r.t 的长期故事是什么？加速和修复的 LLVM 后端？一个会贬低另一个吗？ LLVM IR 的表达能力是否足以实现类型索引数组表示，例如延迟数组？还是计划将repa 和accelerate-llvm 分开的项目？一旦 LLVM 后端成熟，我想它的一个明显基准是 repa？

Accelerate 也使用延迟数组，但它会自动决定使用什么数组表示。这在 Accelerate 中比在 Repa 中更容易，因为 Accelerate 是一种限制性更强的语言，所以我希望我们可以不用 Repa 中需要的表示的类型级提示。我们没有继续开发加速repa 后端的计划。在公认的非常有限的一组基准测试中，新的 LLVM 后端已经比普通的 Repa 代码更快。

【参考方案2】：

一年零几个月前，我在设计yarr 时考虑过这个问题。那时，类型族推断或类似的东西（我不记得确切）存在严重问题，这阻止了实现vector、repa、yarr、accelerate 等的统一包装器。高效并允许不要写太多显式类型签名，或者原则上实现它（我不记得了）。

那是 GHC 7.6。我不知道 GHC 7.8 在这个领域是否有有意义的改进。从理论上讲，我没有看到任何问题，因此我们可以期待这样的东西有一天，无论是短期还是长期，当 GHC 准备好时。