{-# OPTIONS -fno-warn-missing-methods #-} module Data.Array.Parallel.Lifted.Closure ( (:->)(..), PArray(..), mkClosure, mkClosureP, ($:), ($:^), closure, liftedClosure, liftedApply, closure1, closure2, closure3, closure4 ) where import Data.Array.Parallel.PArray.PReprInstances () import Data.Array.Parallel.PArray.PDataInstances import Data.Array.Parallel.Lifted.PArray import GHC.Exts (Int#) infixr 0 :-> infixl 0 $:, $:^ -- | The type of closures. -- This bundles up: -- 1) the vectorised version of the function that takes an explicit environment -- 2) the lifted version, that works on arrays. -- the first parameter of this function is the 'lifting context' -- that gives the length of the array. -- 3) the environment of the closure. -- -- The vectoriser closure-converts the source program so that all functions -- types are expressed in this form. -- data a :-> b = forall e. PA e => Clo !(e -> a -> b) -- vectorised function !(Int# -> PData e -> PData a -> PData b) -- lifted function e -- environment -- | Apply a lifted function by wrapping up the provided array data -- into some real `PArray`s, and passing it those. lifted :: (PArray e -> PArray a -> PArray b) -- ^ lifted function to call. -> Int# -- ^ lifting context -> PData e -- ^ environments -> PData a -- ^ arguments -> PData b -- ^ returned elements {-# INLINE lifted #-} lifted f n# es as = case f (PArray n# es) (PArray n# as) of PArray _ bs -> bs -- | Construct a closure. mkClosure :: forall a b e . PA e => (e -> a -> b) -- ^ vectorised function, with explicit environment. -> (PArray e -> PArray a -> PArray b) -- ^ lifted function, taking an array of environments. -> e -- ^ environment -> (a :-> b) {-# INLINE CONLIKE mkClosure #-} mkClosure fv fl e = Clo fv (lifted fl) e -- | Construct a closure. -- This is like the `mkClosure` function above, except that the provided -- lifted version of the function can take raw array data, instead of -- data wrapped up into a `PArray`. closure :: forall a b e . PA e => (e -> a -> b) -- ^ vectorised function, with explicit environment. -> (Int# -> PData e -> PData a -> PData b) -- ^ lifted function, taking an array of environments. -> e -- ^ environment -> (a :-> b) {-# INLINE closure #-} closure fv fl e = Clo fv fl e -- | Apply a closure to its argument. -- ($:) :: forall a b. (a :-> b) -> a -> b {-# INLINE ($:) #-} Clo f _ e $: a = f e a {-# RULES "mkClosure/($:)" forall fv fl e x. mkClosure fv fl e $: x = fv e x #-} -- | Arrays of closures (aka array closures) -- We need to represent arrays of closures when vectorising partial applications. -- -- For example, consider: -- @mapP (+) xs :: [: Int -> Int :]@ -- -- Representing this an array of thunks doesn't work because we can't evaluate -- in a data parallel manner. Instead, we want *one* function applied to many -- array elements. -- -- Instead, such an array of closures is represented as the vectorised -- and lifted versions of (+), along with an environment array xs that -- contains the partially applied arguments. -- -- @mapP (+) xs ==> AClo plus_v plus_l xs@ -- -- When we find out what the final argument is, we can then use the lifted -- closure application function to compute the result: -- -- @PArray n (AClo plus_v plus_l xs) $:^ (PArray n' ys) -- => PArray n (plus_l n xs ys)@ -- data instance PData (a :-> b) = forall e. PA e => AClo !(e -> a -> b) -- vectorised function, with explicit environment. !(Int# -> PData e -> PData a -> PData b) -- lifted function, taking an array of environments. (PData e) -- array of environments. -- |Lifted closure construction -- mkClosureP :: forall a b e. PA e => (e -> a -> b) -> (PArray e -> PArray a -> PArray b) -> PArray e -> PArray (a :-> b) {-# INLINE mkClosureP #-} mkClosureP fv fl (PArray n# es) = PArray n# (AClo fv (lifted fl) es) liftedClosure :: forall a b e. PA e => (e -> a -> b) -> (Int# -> PData e -> PData a -> PData b) -> PData e -> PData (a :-> b) {-# INLINE liftedClosure #-} liftedClosure fv fl es = AClo fv fl es -- |Lifted closure application -- ($:^) :: forall a b. PArray (a :-> b) -> PArray a -> PArray b {-# INLINE ($:^) #-} PArray n# (AClo _ f es) $:^ PArray _ as = PArray n# (f n# es as) liftedApply :: forall a b. Int# -> PData (a :-> b) -> PData a -> PData b {-# INLINE liftedApply #-} liftedApply n# (AClo _ f es) as = f n# es as -- PRepr instance for closures ------------------------------------------------ type instance PRepr (a :-> b) = a :-> b instance (PA a, PA b) => PA (a :-> b) where toPRepr = id fromPRepr = id toArrPRepr = id fromArrPRepr = id instance PR (a :-> b) where {-# INLINE emptyPR #-} emptyPR = AClo (\_ _ -> error "empty array closure") (\_ _ -> error "empty array closure") (emptyPD :: PData ()) {-# INLINE replicatePR #-} replicatePR n# (Clo f f' e) = AClo f f' (replicatePD n# e) {-# INLINE replicatelPR #-} replicatelPR segd (AClo f f' es) = AClo f f' (replicatelPD segd es) {-# INLINE indexPR #-} indexPR (AClo f f' es) i# = Clo f f' (indexPD es i#) {-# INLINE bpermutePR #-} bpermutePR (AClo f f' es) n# is = AClo f f' (bpermutePD es n# is) {-# INLINE packByTagPR #-} packByTagPR (AClo f f' es) n# tags t# = AClo f f' (packByTagPD es n# tags t#) -- Closure construction ------------------------------------------------------- -- | Arity-1 closures. closure1 :: (a -> b) -> (PArray a -> PArray b) -> (a :-> b) {-# INLINE closure1 #-} closure1 fv fl = mkClosure (\_ -> fv) (\_ -> fl) () -- | Arity-2 closures. closure2 :: PA a => (a -> b -> c) -> (PArray a -> PArray b -> PArray c) -> (a :-> b :-> c) {-# INLINE closure2 #-} closure2 fv fl = mkClosure fv_1 fl_1 () where fv_1 _ x = mkClosure fv fl x fl_1 _ xs = mkClosureP fv fl xs -- | Arity-3 closures. closure3 :: (PA a, PA b) => (a -> b -> c -> d) -> (PArray a -> PArray b -> PArray c -> PArray d) -> (a :-> b :-> c :-> d) {-# INLINE closure3 #-} closure3 fv fl = mkClosure fv_1 fl_1 () where fv_1 _ x = mkClosure fv_2 fl_2 x fl_1 _ xs = mkClosureP fv_2 fl_2 xs fv_2 x y = mkClosure fv_3 fl_3 (x,y) fl_2 xs ys = mkClosureP fv_3 fl_3 (zipPA# xs ys) fv_3 (x,y) z = fv x y z fl_3 ps zs = case unzipPA# ps of (xs,ys) -> fl xs ys zs -- | Arity-4 closures. closure4 :: (PA a, PA b, PA c) => (a -> b -> c -> d -> e) -> (PArray a -> PArray b -> PArray c -> PArray d -> PArray e) -> (a :-> b :-> c :-> d :-> e) {-# INLINE closure4 #-} closure4 fv fl = mkClosure fv_1 fl_1 () where fv_1 _ x = mkClosure fv_2 fl_2 x fl_1 _ xs = mkClosureP fv_2 fl_2 xs fv_2 x y = mkClosure fv_3 fl_3 (x, y) fl_2 xs ys = mkClosureP fv_3 fl_3 (zipPA# xs ys) fv_3 (x, y) z = mkClosure fv_4 fl_4 (x, y, z) fl_3 xys zs = case unzipPA# xys of (xs, ys) -> mkClosureP fv_4 fl_4 (zip3PA# xs ys zs) fv_4 (x, y, z) v = fv x y z v fl_4 ps vs = case unzip3PA# ps of (xs, ys, zs) -> fl xs ys zs vs