module Data.Array.Accelerate.Utility.Loop (
nest,
nestLog2,
) where
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate (Acc, Exp)
nest, _nest ::
(A.Arrays a) =>
Exp Int -> (Acc a -> Acc a) -> Acc a -> Acc a
nest n0 f x0 =
A.asnd $
A.awhile
(A.map (A.> 0) . A.afst)
(A.lift . (\(n, x) -> (A.map (subtract 1) n, f x)) . A.unlift)
(A.lift (A.unit n0, x0))
_nest n0 f x0 =
A.asnd $
A.awhile
(A.unit . (A.> 0) . A.the . A.afst)
(A.lift . (\(n, x) -> (A.unit $ A.the n - 1, f x)) . A.unlift)
(A.lift (A.unit n0, x0))
nestLog2 ::
(A.Arrays a) =>
Exp Int -> (Acc a -> Acc a) -> Acc a -> Acc a
nestLog2 n0 f x0 =
A.asnd $
A.awhile
(A.map (A.> 1) . A.afst)
(A.lift . (\(n, x) -> (A.map (flip div 2 . (1+)) n, f x)) . A.unlift)
(A.lift (A.unit n0, x0))