Zenn
🌊

Haskell-Accelerateでのループ

2025/01/01に公開
Haskell
並列処理
GPGPU
accelerate
tech

背景

多用するけどオリジナルのリファレンスには書かれていない用法だったので書いておきます。
ループというか指定した回数だけ関数を重ねがけするものです。GPU環境での実験ですが、CPU実行でも少量変更で良いはずです。

環境

名前 バージョン
OS Ubuntu 20.04
nvcc(cuda) 11.8
ghc 8.6.5
accelerate 1.3.0.0
accelerate-llvm 1.3.0.0
accelerate-llvm-native 1.3.0.0
accelerate-llvm-ptx 1.3.0.0
GPU GeForce RTX3080
CPU AMD Ryzen 9 3900XT

(>->)によるループ

実装としては簡単なのですが、Accのネストが深くなり速度が出ません。
以下はn要素のベクトルの先頭5要素のみを倍にする操作をn回かけるものです。
fCalc = L.foldl1 (>->) $ L.replicate n fiveDoubleによりfiveDoubleをn個結合して1つのAcc関数を構成しています。

loop_compose.hs
{-# LANGUAGE CPP, BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}

import Data.Array.Accelerate as A
import Data.Array.Accelerate.LLVM.PTX
import Data.List as L
import Control.DeepSeq (deepseq)

fiveDouble :: Acc (Vector Int) -> Acc (Vector Int)
fiveDouble v = generate (shape v) f
  where
    f p@(I1 i) = ifThenElse (i A.== -1) orgVal -- dummy
                 $ ifThenElse (i A.== -1) orgVal -- dummy
                 $ ifThenElse (i A.< 5) (orgVal * 2) orgVal
                     where orgVal = v ! p
               
nop :: Acc (Vector Int) -> Acc (Scalar Int)
nop _ = fill I0 0

main = do
  let
    n = 10
    -- n = 10000
    v = enumFromN (constant (Z :. n)) 0 :: Acc(Vector Int)
  
    fCalc :: Acc (Vector Int) -> Acc (Vector Int)
    fCalc = L.foldl1 (>->) $ L.replicate n fiveDouble
      
  print $ run $ v
  print $ run $ fiveDouble v
  print $ run $ fCalc v  
  print $ fCalc v -- to show Acc tree.
  -- (run $ (fCalc >-> nop) v) `deepseq` return ()
実行結果
Vector (Z :. 10) [0,1,2,3,4,5,6,7,8,9]
Vector (Z :. 10) [0,2,4,6,8,5,6,7,8,9]
Vector (Z :. 10) [0,1024,2048,3072,4096,5,6,7,8,9]
let
  a0 = (\a0 ->
          let
            a1 = (\a1 ->
                    let
                      a2 = (\a2 ->
                              let
                                a3 = (\a3 ->
                                        let
                                          a4 = (\a4 ->
                                                  let
                                                    a5 = (\a5 ->
                                                            let
                                                              a6 = (\a6 ->
                                                                      let
                                                                        a7 = (\a7 ->
                                                                                let
                                                                                  a8 = (\a8 ->
                                                                                          generate
                                                                                            (shape a8)
                                                                                            (\(T1 x0) ->
                                                                                               let
                                                                                                 x1 = a8 ! x0
                                                                                                 x2 = -1 == x0
                                                                                               in
                                                                                               if x2
                                                                                                  then x1
省略

awhileによるループ

少々手間ですがAccのネストが浅く速度も出せます。
awhile にて、メインで扱いたいVectorとセットでカウンタの役割をするScalarをペアにAcc (Scalar Int, Vector Int)として扱います。カウンタに基づいたVector編集などにも使えるので便利でしょう。
注意： 普通のfor分と同じノリで使うと危ないです。カウンタが終了条件にかかる直前の処理でもメインの処理も実行してしまうので異常なメモリアクセスをする場合があり、扱うデータサイズが大きくなってからメモリアクセスエラーとなります。以下もやってるかも。後々修正します。

loop_awhile.hs
{-# LANGUAGE CPP, BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}

import Data.Array.Accelerate as A
import Data.Array.Accelerate.LLVM.PTX
import Data.List as L
import Control.DeepSeq (deepseq)

fiveDouble :: Acc (Vector Int) -> Acc (Vector Int)
fiveDouble v = generate (shape v) f
  where
    f p@(I1 i) = ifThenElse (i A.== -1) orgVal -- dummy
                 $ ifThenElse (i A.== -1) orgVal -- dummy
                 $ ifThenElse (i A.< 5) (orgVal * 2) orgVal
                     where orgVal = v ! p
               
nop :: Acc (Vector Int) -> Acc (Scalar Int)
nop _ = fill I0 0

main = do
  let
    n = 10
    -- n = 10000
    v = enumFromN (constant (Z :. n)) 0 :: Acc(Vector Int)
  
    fCalc :: Acc (Vector Int) -> Acc (Vector Int)
    fCalc xs = asnd $ awhile cond f ini
      where
        ini = T2 (unit 0) xs :: Acc (Scalar Int, Vector Int)
  
        f :: Acc (Scalar Int, Vector Int) -> Acc (Scalar Int, Vector Int)
        f (T2 c xs) = T2 (unit $ (the c)+1) (fiveDouble xs)
  
        cond :: Acc (Scalar Int, Vector Int) -> Acc (Scalar Bool)
        cond (T2 c _) = unit $ (the c) A.< (lift n)

  print $ run $ v
  print $ run $ fiveDouble v  
  print $ run $ fCalc v  
  print $ fCalc v -- to show Acc tree.
  -- (run $ (fCalc >-> nop) v) `deepseq` return ()
実行結果
Vector (Z :. 10) [0,1,2,3,4,5,6,7,8,9]
Vector (Z :. 10) [0,2,4,6,8,5,6,7,8,9]
Vector (Z :. 10) [0,1024,2048,3072,4096,5,6,7,8,9]
let
  T2 a0 a1 = awhile
               (\(T2 a0 a1) -> let a2 = unit (a0 ! () < 10, ()) in map (\(x0, ()) -> x0) a2)
               (\(T2 a0 a1) ->
                  T2 (unit (1 + a0 ! ()))
                  (generate
                     (shape a1)
                     (\(T1 x0) -> let x1 = a1 ! x0 x2 = -1 == x0 in x2 ? x1 : x2 ? x1 : x0 < 5 ? 2 * x1 : x1)))
               T2 (unit 0) (generate 10 (\(T1 x0) -> x0))
in
a1

比較

n=10000として以下の様に変更しtimeコマンドで速度比較しました。nopはGPUからのデータ転送の時間を除くためのものです。

loop_compose.hs
{-# LANGUAGE CPP, BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}

import Data.Array.Accelerate as A
import Data.Array.Accelerate.LLVM.PTX
import Data.List as L
import Control.DeepSeq (deepseq)

fiveDouble :: Acc (Vector Int) -> Acc (Vector Int)
fiveDouble v = generate (shape v) f
  where
    f p@(I1 i) = ifThenElse (i A.== -1) orgVal -- dummy
                 $ ifThenElse (i A.== -1) orgVal -- dummy
                 $ ifThenElse (i A.< 5) (orgVal * 2) orgVal
                     where orgVal = v ! p
               
nop :: Acc (Vector Int) -> Acc (Scalar Int)
nop _ = fill I0 0

main = do
  let
    -- n = 10
    n = 10000
    v = enumFromN (constant (Z :. n)) 0 :: Acc(Vector Int)
  
    fCalc :: Acc (Vector Int) -> Acc (Vector Int)
    fCalc = L.foldl1 (>->) $ L.replicate n fiveDouble
      
  --print $ run $ v
  --print $ run $ fiveDouble v  
  --print $ run $ fCalc v  
  --print $ fCalc v -- to show Acc tree.
  (run $ (fCalc >-> nop) v) `deepseq` return ()
loop_awhile.hs
{-# LANGUAGE CPP, BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}

import Data.Array.Accelerate as A
import Data.Array.Accelerate.LLVM.PTX
import Data.List as L
import Control.DeepSeq (deepseq)

fiveDouble :: Acc (Vector Int) -> Acc (Vector Int)
fiveDouble v = generate (shape v) f
  where
    f p@(I1 i) = ifThenElse (i A.== -1) orgVal -- dummy
                 $ ifThenElse (i A.== -1) orgVal -- dummy
                 $ ifThenElse (i A.< 5) (orgVal * 2) orgVal
                     where orgVal = v ! p
               
nop :: Acc (Vector Int) -> Acc (Scalar Int)
nop _ = fill I0 0

main = do
  let
    -- n = 10
    n = 10000
    v = enumFromN (constant (Z :. n)) 0 :: Acc(Vector Int)
  
    fCalc :: Acc (Vector Int) -> Acc (Vector Int)
    fCalc xs = asnd $ awhile cond f ini
      where
        ini = T2 (unit 0) xs :: Acc (Scalar Int, Vector Int)
  
        f :: Acc (Scalar Int, Vector Int) -> Acc (Scalar Int, Vector Int)
        f (T2 c xs) = T2 (unit $ (the c)+1) (fiveDouble xs)
  
        cond :: Acc (Scalar Int, Vector Int) -> Acc (Scalar Bool)
        cond (T2 c _) = unit $ (the c) A.< (lift n)

  -- print $ run $ v
  -- print $ run $ fiveDouble v  
  -- print $ run $ fCalc v  
  -- print $ fCalc v -- to show Acc tree.
  (run $ (fCalc >-> nop) v) `deepseq` return ()

結果：

$ghc -O2 loop_compose -threaded
$ghc -O2 loop_awhile -threaded
$ time ./loop_compose 
real	0m8.689s
user	0m7.787s
sys	0m0.852s

$ time ./loop_awhile 
real	0m1.889s
user	0m1.631s
sys	0m0.247s

4.6倍awhileが早い結果となりました。
GPUメモリ消費量はどちらも1Gbyte程度でした。すべてのケースで>->が遅いわけではないです。後々早くなる例を掲載します。

参考文献

>->の方法:https://www.oreilly.com/library/view/parallel-and-concurrent/9781449335939/ch06.html
awhileの方法:https://github.com/AccelerateHS/accelerate-examples/blob/master/examples/quicksort/QuickSort.hs

Discussion