module Futhark.Optimise.Fusion.Screma
  ( splitLambdaByPar,
    fuseScrema,
    fuseSuperScrema,
    SuperScrema (..),
    moveRedScanSuperScrema,
    moveLastSuperScrema,
    fusible,
    toScrema,
  )
where

import Control.Monad
import Data.Bifunctor
import Data.Function
import Data.List qualified as L
import Data.Map.Strict qualified as M
import Data.Maybe
import Futhark.Analysis.Alias
import Futhark.Analysis.DataDependencies
import Futhark.Analysis.HORep.SOAC qualified as SOAC
import Futhark.IR
import Futhark.IR.Prop.Aliases
import Futhark.IR.SOACS
import Futhark.IR.SOACS.Simplify
import Futhark.MonadFreshNames
import Futhark.Tools
import Futhark.Transform.Rename
import Futhark.Util (splitAt3)
import Futhark.Util.Pretty

-- | The fused representation of two scremas.
data SuperScrema rep
  = SuperScrema
      SubExp
      [SOAC.Input]
      (Lambda rep)
      [Scan rep]
      [Reduce rep]
      (Lambda rep)
      [Scan rep]
      [Reduce rep]
      (Lambda rep)
  deriving (SuperScrema rep -> SuperScrema rep -> Bool
(SuperScrema rep -> SuperScrema rep -> Bool)
-> (SuperScrema rep -> SuperScrema rep -> Bool)
-> Eq (SuperScrema rep)
forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Bool
== :: SuperScrema rep -> SuperScrema rep -> Bool
$c/= :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Bool
/= :: SuperScrema rep -> SuperScrema rep -> Bool
Eq, Eq (SuperScrema rep)
Eq (SuperScrema rep) =>
(SuperScrema rep -> SuperScrema rep -> Ordering)
-> (SuperScrema rep -> SuperScrema rep -> Bool)
-> (SuperScrema rep -> SuperScrema rep -> Bool)
-> (SuperScrema rep -> SuperScrema rep -> Bool)
-> (SuperScrema rep -> SuperScrema rep -> Bool)
-> (SuperScrema rep -> SuperScrema rep -> SuperScrema rep)
-> (SuperScrema rep -> SuperScrema rep -> SuperScrema rep)
-> Ord (SuperScrema rep)
SuperScrema rep -> SuperScrema rep -> Bool
SuperScrema rep -> SuperScrema rep -> Ordering
SuperScrema rep -> SuperScrema rep -> SuperScrema rep
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall rep. RepTypes rep => Eq (SuperScrema rep)
forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Bool
forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Ordering
forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> SuperScrema rep
$ccompare :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Ordering
compare :: SuperScrema rep -> SuperScrema rep -> Ordering
$c< :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Bool
< :: SuperScrema rep -> SuperScrema rep -> Bool
$c<= :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Bool
<= :: SuperScrema rep -> SuperScrema rep -> Bool
$c> :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Bool
> :: SuperScrema rep -> SuperScrema rep -> Bool
$c>= :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> Bool
>= :: SuperScrema rep -> SuperScrema rep -> Bool
$cmax :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> SuperScrema rep
max :: SuperScrema rep -> SuperScrema rep -> SuperScrema rep
$cmin :: forall rep.
RepTypes rep =>
SuperScrema rep -> SuperScrema rep -> SuperScrema rep
min :: SuperScrema rep -> SuperScrema rep -> SuperScrema rep
Ord, Int -> SuperScrema rep -> ShowS
[SuperScrema rep] -> ShowS
SuperScrema rep -> String
(Int -> SuperScrema rep -> ShowS)
-> (SuperScrema rep -> String)
-> ([SuperScrema rep] -> ShowS)
-> Show (SuperScrema rep)
forall rep. RepTypes rep => Int -> SuperScrema rep -> ShowS
forall rep. RepTypes rep => [SuperScrema rep] -> ShowS
forall rep. RepTypes rep => SuperScrema rep -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall rep. RepTypes rep => Int -> SuperScrema rep -> ShowS
showsPrec :: Int -> SuperScrema rep -> ShowS
$cshow :: forall rep. RepTypes rep => SuperScrema rep -> String
show :: SuperScrema rep -> String
$cshowList :: forall rep. RepTypes rep => [SuperScrema rep] -> ShowS
showList :: [SuperScrema rep] -> ShowS
Show)

instance (PrettyRep rep) => Pretty (SuperScrema rep) where
  pretty :: forall ann. SuperScrema rep -> Doc ann
pretty (SuperScrema SubExp
w [Input]
inps Lambda rep
lam1 [Scan rep]
scans1 [Reduce rep]
reds1 Lambda rep
lam2 [Scan rep]
scans2 [Reduce rep]
reds2 Lambda rep
lam3) =
    Doc ann
"superscrema"
      Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> (Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align)
        ( (SubExp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. SubExp -> Doc ann
pretty SubExp
w Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma)
            Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
ppTuple' ((Input -> Doc ann) -> [Input] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map Input -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Input -> Doc ann
pretty [Input]
inps)
            Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Lambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Lambda rep -> Doc ann
pretty Lambda rep
lam1
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
              Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> [Scan rep] -> Doc ann
forall {a} {ann}. Pretty a => [a] -> Doc ann
p' [Scan rep]
scans1
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
              Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> [Reduce rep] -> Doc ann
forall {a} {ann}. Pretty a => [a] -> Doc ann
p' [Reduce rep]
reds1
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
              Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Lambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Lambda rep -> Doc ann
pretty Lambda rep
lam2
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
              Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> [Scan rep] -> Doc ann
forall {a} {ann}. Pretty a => [a] -> Doc ann
p' [Scan rep]
scans2
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
              Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> [Reduce rep] -> Doc ann
forall {a} {ann}. Pretty a => [a] -> Doc ann
p' [Reduce rep]
reds2
            Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
comma
              Doc ann -> Doc ann -> Doc ann
forall a. Doc a -> Doc a -> Doc a
</> Lambda rep -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Lambda rep -> Doc ann
pretty Lambda rep
lam3
        )
    where
      p' :: [a] -> Doc ann
p' [a]
xs = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
braces ([Doc ann] -> Doc ann
forall a. Monoid a => [a] -> a
mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
L.intersperse (Doc ann
forall ann. Doc ann
comma Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
forall ann. Doc ann
line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
forall a b. (a -> b) -> a -> b
$ (a -> Doc ann) -> [a] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty [a]
xs)

-- | Using a boolean mask pick elements in the
pick :: [Bool] -> [a] -> [a]
pick :: forall a. [Bool] -> [a] -> [a]
pick [Bool]
bs [a]
xs = ((Bool, a) -> a) -> [(Bool, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Bool, a) -> a
forall a b. (a, b) -> b
snd ([(Bool, a)] -> [a]) -> [(Bool, a)] -> [a]
forall a b. (a -> b) -> a -> b
$ ((Bool, a) -> Bool) -> [(Bool, a)] -> [(Bool, a)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool, a) -> Bool
forall a b. (a, b) -> a
fst ([(Bool, a)] -> [(Bool, a)]) -> [(Bool, a)] -> [(Bool, a)]
forall a b. (a -> b) -> a -> b
$ [Bool] -> [a] -> [(Bool, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Bool]
bs [a]
xs

-- | Check that all inputs and outputs are fusible for a producer and
-- consumer.
fuseIsVarish :: [SOAC.Input] -> [VName] -> Bool
fuseIsVarish :: [Input] -> [VName] -> Bool
fuseIsVarish [Input]
inp_c =
  (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> (Input -> Bool) -> Maybe Input -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Maybe VName -> Bool
forall a. Maybe a -> Bool
isJust (Maybe VName -> Bool) -> (Input -> Maybe VName) -> Input -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> Maybe VName
SOAC.isVarishInput) (Maybe Input -> Bool) -> (VName -> Maybe Input) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Map VName Input -> Maybe Input)
-> Map VName Input -> VName -> Maybe Input
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> Map VName Input -> Maybe Input
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Map VName Input
name_to_inp_c)
  where
    name_to_inp_c :: Map VName Input
name_to_inp_c = [(VName, Input)] -> Map VName Input
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Input)] -> Map VName Input)
-> [(VName, Input)] -> Map VName Input
forall a b. (a -> b) -> a -> b
$ [VName] -> [Input] -> [(VName, Input)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Input -> VName
SOAC.inputArray (Input -> VName) -> [Input] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Input]
inp_c) [Input]
inp_c

-- | Given a list of parameter names and a lambda, split the lambda
-- into two where the first function will depend on the given
-- parameters given if they exist plus additional parameters if they
-- are need for some computation. The other lambda will have all the
-- other parameters and do the remaining computational work. Some
-- statements may be present in both lambdas.
--
-- Further, each parameter is associated with some additional
-- information @a@, and each result with some additional information
-- @b@. This is also partitioned and returned appropriately.
splitLambdaByPar ::
  (MonadFail m) =>
  [VName] ->
  [inp] ->
  Lambda SOACS ->
  [out] ->
  m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
splitLambdaByPar :: forall (m :: * -> *) inp out.
MonadFail m =>
[VName]
-> [inp]
-> Lambda SOACS
-> [out]
-> m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
splitLambdaByPar [VName]
names [inp]
inp Lambda SOACS
lam [out]
out = do
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when
    (Lambda SOACS -> Lambda SOACS -> Bool
consumedOverlap Lambda SOACS
new_lam Lambda SOACS
new_lam')
    (String -> m ()
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Can not fuse due to overlap consumption in pre and post.")
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when
    (Lambda SOACS -> Lambda SOACS -> Bool
parAccsOverlap Lambda SOACS
new_lam Lambda SOACS
new_lam')
    (String -> m ()
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Can not fuse due to overlap in parameter accumalators.")
  (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
-> m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( ([inp]
new_inp, Lambda SOACS
new_lam, [out]
new_out),
      ([inp]
new_inp', Lambda SOACS
new_lam', [out]
new_out')
    )
  where
    new_lam :: Lambda SOACS
new_lam = Lambda SOACS -> Lambda SOACS
forall rep. Buildable rep => Lambda rep -> Lambda rep
eliminateByRes (Lambda SOACS -> Lambda SOACS) -> Lambda SOACS -> Lambda SOACS
forall a b. (a -> b) -> a -> b
$ [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
new_params [Type]
new_ts (Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
stms [SubExpRes]
new_res)
    new_lam' :: Lambda SOACS
new_lam' = Lambda SOACS -> Lambda SOACS
forall rep. Buildable rep => Lambda rep -> Lambda rep
eliminateByRes (Lambda SOACS -> Lambda SOACS) -> Lambda SOACS -> Lambda SOACS
forall a b. (a -> b) -> a -> b
$ [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
new_params' [Type]
new_ts' (Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
stms [SubExpRes]
new_res')
    pars :: [LParam SOACS]
pars = Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam
    m :: Map (Param Type) inp
m = [(Param Type, inp)] -> Map (Param Type) inp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Param Type, inp)] -> Map (Param Type) inp)
-> [(Param Type, inp)] -> Map (Param Type) inp
forall a b. (a -> b) -> a -> b
$ [Param Type] -> [inp] -> [(Param Type, inp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
pars [inp]
inp
    par_deps :: [Names]
par_deps = AliasTable -> Lambda SOACS -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies AliasTable
forall a. Monoid a => a
mempty Lambda SOACS
lam (VName -> Names
oneName (VName -> Names) -> (Param Type -> VName) -> Param Type -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> Names) -> [Param Type] -> [Names]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Param Type]
pars)
    body :: Body SOACS
body = Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam
    stms :: Stms SOACS
stms = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms Body SOACS
body
    new_inp :: [inp]
new_inp = (Map (Param Type) inp
m Map (Param Type) inp -> Param Type -> inp
forall k a. Ord k => Map k a -> k -> a
M.!) (Param Type -> inp) -> [Param Type] -> [inp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Param Type]
new_params
    new_inp' :: [inp]
new_inp' = (Map (Param Type) inp
m Map (Param Type) inp -> Param Type -> inp
forall k a. Ord k => Map k a -> k -> a
M.!) (Param Type -> inp) -> [Param Type] -> [inp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Param Type]
new_params'
    new_params :: [Param Type]
new_params = (Param Type -> Bool) -> [Param Type] -> [Param Type]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> Names -> Bool
`nameIn` Names
deps) (VName -> Bool) -> (Param Type -> VName) -> Param Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName) [Param Type]
pars
    new_params' :: [Param Type]
new_params' = (Param Type -> Bool) -> [Param Type] -> [Param Type]
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> Names -> Bool
`nameIn` Names
deps') (VName -> Bool) -> (Param Type -> VName) -> Param Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName) [Param Type]
pars
    auxiliary :: [(Names, b, c, d)] -> (Names, [b], [c], [d])
auxiliary = (\([Names]
a, [b]
b, [c]
c, [d]
d) -> ([Names] -> Names
forall a. Monoid a => [a] -> a
mconcat [Names]
a, [b]
b, [c]
c, [d]
d)) (([Names], [b], [c], [d]) -> (Names, [b], [c], [d]))
-> ([(Names, b, c, d)] -> ([Names], [b], [c], [d]))
-> [(Names, b, c, d)]
-> (Names, [b], [c], [d])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Names, b, c, d)] -> ([Names], [b], [c], [d])
forall a b c d. [(a, b, c, d)] -> ([a], [b], [c], [d])
L.unzip4
    ((Names
deps, [SubExpRes]
new_res, [Type]
new_ts, [out]
new_out), (Names
deps', [SubExpRes]
new_res', [Type]
new_ts', [out]
new_out')) =
      ([(Names, SubExpRes, Type, out)]
 -> (Names, [SubExpRes], [Type], [out]))
-> ([(Names, SubExpRes, Type, out)]
    -> (Names, [SubExpRes], [Type], [out]))
-> ([(Names, SubExpRes, Type, out)],
    [(Names, SubExpRes, Type, out)])
-> ((Names, [SubExpRes], [Type], [out]),
    (Names, [SubExpRes], [Type], [out]))
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap [(Names, SubExpRes, Type, out)]
-> (Names, [SubExpRes], [Type], [out])
forall {b} {c} {d}. [(Names, b, c, d)] -> (Names, [b], [c], [d])
auxiliary [(Names, SubExpRes, Type, out)]
-> (Names, [SubExpRes], [Type], [out])
forall {b} {c} {d}. [(Names, b, c, d)] -> (Names, [b], [c], [d])
auxiliary
        (([(Names, SubExpRes, Type, out)], [(Names, SubExpRes, Type, out)])
 -> ((Names, [SubExpRes], [Type], [out]),
     (Names, [SubExpRes], [Type], [out])))
-> ([(Names, SubExpRes, Type, out)]
    -> ([(Names, SubExpRes, Type, out)],
        [(Names, SubExpRes, Type, out)]))
-> [(Names, SubExpRes, Type, out)]
-> ((Names, [SubExpRes], [Type], [out]),
    (Names, [SubExpRes], [Type], [out]))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Names, SubExpRes, Type, out) -> Bool)
-> [(Names, SubExpRes, Type, out)]
-> ([(Names, SubExpRes, Type, out)],
    [(Names, SubExpRes, Type, out)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition (Names -> Names -> Bool
namesIntersect ([VName] -> Names
namesFromList [VName]
names) (Names -> Bool)
-> ((Names, SubExpRes, Type, out) -> Names)
-> (Names, SubExpRes, Type, out)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\(Names
a, SubExpRes
_, Type
_, out
_) -> Names
a))
        ([(Names, SubExpRes, Type, out)]
 -> ((Names, [SubExpRes], [Type], [out]),
     (Names, [SubExpRes], [Type], [out])))
-> [(Names, SubExpRes, Type, out)]
-> ((Names, [SubExpRes], [Type], [out]),
    (Names, [SubExpRes], [Type], [out]))
forall a b. (a -> b) -> a -> b
$ [Names]
-> [SubExpRes]
-> [Type]
-> [out]
-> [(Names, SubExpRes, Type, out)]
forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
L.zip4 [Names]
par_deps (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult Body SOACS
body) (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam) [out]
out

-- | Check that two scremas are fusible if they are give back the
-- producer scremas post lambda that has been split into the scan
-- lambda and map lambda. It is not fusible if inputs and outputs are
-- being transformed and if the producers scan result is used for the
-- consumers scans or reduces.
fusible ::
  (MonadFail m) =>
  ScremaForm SOACS ->
  [VName] ->
  [SOAC.Input] ->
  ScremaForm SOACS ->
  m ()
fusible :: forall (m :: * -> *).
MonadFail m =>
ScremaForm SOACS -> [VName] -> [Input] -> ScremaForm SOACS -> m ()
fusible ScremaForm SOACS
form_p [VName]
out_p [Input]
inp_c ScremaForm SOACS
form_c = do
  ((_, post_scan_p, _), _) <-
    [VName]
-> [Param Type]
-> Lambda SOACS
-> [Type]
-> m (([Param Type], Lambda SOACS, [Type]),
      ([Param Type], Lambda SOACS, [Type]))
forall (m :: * -> *) inp out.
MonadFail m =>
[VName]
-> [inp]
-> Lambda SOACS
-> [out]
-> m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
splitLambdaByPar [VName]
post_scan_pars_p (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
post_p) Lambda SOACS
post_p (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
post_p)
  let post_scan_res_p = Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
post_scan_p
      forbidden_p = [VName] -> Names
namesFromList ([VName] -> Names) -> [VName] -> Names
forall a b. (a -> b) -> a -> b
$ [VName] -> Lambda SOACS -> SubExpRes -> VName
resToOut [VName]
out_p Lambda SOACS
post_p (SubExpRes -> VName) -> [SubExpRes] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SubExpRes]
post_scan_res_p
      is_fusible =
        [Input] -> [VName] -> Bool
fuseIsVarish [Input]
inp_c [VName]
out_p
          Bool -> Bool -> Bool
&& Bool -> Bool
not (Names
forbidden_c Names -> Names -> Bool
`namesIntersect` Names
forbidden_p)
          Bool -> Bool -> Bool
&& (VName -> Bool) -> [VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all
            (VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` (Input -> Maybe VName) -> [Input] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Input -> Maybe VName
SOAC.isVarishInput [Input]
inp_c)
            (Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
take Int
num_red_p [VName]
out_p)
  unless is_fusible (fail "Scremas are not fusible.")
  where
    pre_pars_c :: [Names]
pre_pars_c = VName -> Names
oneName (VName -> Names) -> (Param Type -> VName) -> Param Type -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> Names) -> [Param Type] -> [Names]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
pre_c
    ([Names]
pre_scan_deps_c, [Names]
pre_red_deps_c, [Names]
_) =
      Int -> Int -> [Names] -> ([Names], [Names], [Names])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3 Int
num_scan_c Int
num_red_c ([Names] -> ([Names], [Names], [Names]))
-> [Names] -> ([Names], [Names], [Names])
forall a b. (a -> b) -> a -> b
$
        AliasTable -> Lambda SOACS -> [Names] -> [Names]
forall rep t.
(ASTRep rep, FreeIn t) =>
AliasTable -> GLambda rep t -> [Names] -> [Names]
lambdaDependencies AliasTable
forall a. Monoid a => a
mempty Lambda SOACS
pre_c [Names]
pre_pars_c
    forbidden_c :: Names
forbidden_c =
      [VName] -> Names
namesFromList
        ([VName] -> Names) -> (Names -> [VName]) -> Names -> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Maybe VName) -> [VName] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((Input -> VName) -> Maybe Input -> Maybe VName
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Input -> VName
SOAC.inputArray (Maybe Input -> Maybe VName)
-> (VName -> Maybe Input) -> VName -> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Input] -> Lambda SOACS -> VName -> Maybe Input
parToInp [Input]
inp_c Lambda SOACS
pre_c)
        ([VName] -> [VName]) -> (Names -> [VName]) -> Names -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Names -> [VName]
namesToList
        (Names -> Names) -> Names -> Names
forall a b. (a -> b) -> a -> b
$ [Names] -> Names
forall a. Monoid a => [a] -> a
mconcat ([Names]
pre_scan_deps_c [Names] -> [Names] -> [Names]
forall a. Semigroup a => a -> a -> a
<> [Names]
pre_red_deps_c)
    pre_c :: Lambda SOACS
pre_c = ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_c
    post_p :: Lambda SOACS
post_p = ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_p
    post_scan_pars_p :: [VName]
post_scan_pars_p = Int -> [VName] -> [VName]
forall a. Int -> [a] -> [a]
take Int
num_scan_p ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> VName) -> [Param Type] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
post_p
    num_scan_c :: Int
num_scan_c = [Scan SOACS] -> Int
forall rep. [Scan rep] -> Int
scanResults ([Scan SOACS] -> Int) -> [Scan SOACS] -> Int
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> [Scan SOACS]
forall rep. ScremaForm rep -> [Scan rep]
scremaScans ScremaForm SOACS
form_c
    num_red_c :: Int
num_red_c = [Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults ([Reduce SOACS] -> Int) -> [Reduce SOACS] -> Int
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> [Reduce SOACS]
forall rep. ScremaForm rep -> [Reduce rep]
scremaReduces ScremaForm SOACS
form_c
    num_red_p :: Int
num_red_p = [Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults ([Reduce SOACS] -> Int) -> [Reduce SOACS] -> Int
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> [Reduce SOACS]
forall rep. ScremaForm rep -> [Reduce rep]
scremaReduces ScremaForm SOACS
form_p
    num_scan_p :: Int
num_scan_p = [Scan SOACS] -> Int
forall rep. [Scan rep] -> Int
scanResults ([Scan SOACS] -> Int) -> [Scan SOACS] -> Int
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> [Scan SOACS]
forall rep. ScremaForm rep -> [Scan rep]
scremaScans ScremaForm SOACS
form_p

-- | Given two scremas that are fusible, fuse them into a super
-- screma. This is fused but work will have to be moved around in the
-- super screma for it to become a screma.
fuseSuperScrema ::
  (MonadFreshNames m) =>
  SubExp ->
  [SOAC.Input] ->
  ScremaForm SOACS ->
  [VName] ->
  [SOAC.Input] ->
  ScremaForm SOACS ->
  [VName] ->
  m (SuperScrema SOACS, [VName])
fuseSuperScrema :: forall (m :: * -> *).
MonadFreshNames m =>
SubExp
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> m (SuperScrema SOACS, [VName])
fuseSuperScrema SubExp
w [Input]
inp_p' ScremaForm SOACS
form_p' [VName]
out_p [Input]
inp_c' ScremaForm SOACS
form_c' [VName]
out_c = do
  let ([Input]
inp_p, ScremaForm SOACS
form_p) = [Input] -> ScremaForm SOACS -> ([Input], ScremaForm SOACS)
forall rep a.
(Buildable rep, Ord a) =>
[a] -> ScremaForm rep -> ([a], ScremaForm rep)
dedupInput [Input]
inp_p' ScremaForm SOACS
form_p'
      ([Input]
inp_c, ScremaForm SOACS
form_c) = [Input] -> ScremaForm SOACS -> ([Input], ScremaForm SOACS)
forall rep a.
(Buildable rep, Ord a) =>
[a] -> ScremaForm rep -> ([a], ScremaForm rep)
dedupInput [Input]
inp_c' ScremaForm SOACS
form_c'
      inp_c_real_map :: [Bool]
inp_c_real_map = (Input -> Bool) -> [Input] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> Bool
not (Bool -> Bool) -> (Input -> Bool) -> Input -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Input -> Bool
inputFromOutput) [Input]
inp_c
      inp_c_real :: [Input]
inp_c_real = [Bool] -> [Input] -> [Input]
forall a. [Bool] -> [a] -> [a]
pick [Bool]
inp_c_real_map [Input]
inp_c
      inp_r :: [Input]
inp_r = [Input]
inp_p [Input] -> [Input] -> [Input]
forall a. Semigroup a => a -> a -> a
<> [Input]
inp_c_real

      ([VName]
out_red_p, [VName]
out_post_p) =
        Int -> [VName] -> ([VName], [VName])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults ([Reduce SOACS] -> Int) -> [Reduce SOACS] -> Int
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> [Reduce SOACS]
forall rep. ScremaForm rep -> [Reduce rep]
scremaReduces ScremaForm SOACS
form_p) [VName]
out_p
      ([VName]
out_red_c, [VName]
out_post_c) =
        Int -> [VName] -> ([VName], [VName])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults ([Reduce SOACS] -> Int) -> [Reduce SOACS] -> Int
forall a b. (a -> b) -> a -> b
$ ScremaForm SOACS -> [Reduce SOACS]
forall rep. ScremaForm rep -> [Reduce rep]
scremaReduces ScremaForm SOACS
form_c) [VName]
out_c

  forward_params <- [Param Type] -> (Param Type -> m (Param Type)) -> m [Param Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ([Bool] -> [Param Type] -> [Param Type]
forall a. [Bool] -> [a] -> [a]
pick [Bool]
inp_c_real_map (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_c))) ((Param Type -> m (Param Type)) -> m [Param Type])
-> (Param Type -> m (Param Type)) -> m [Param Type]
forall a b. (a -> b) -> a -> b
$ \Param Type
p ->
    Name -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam (VName -> Name
baseName (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
p)) (Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param Type
p)

  let lam1 =
        Lambda
          { lambdaParams :: [LParam SOACS]
lambdaParams =
              Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_p) [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type]
forward_params,
            lambdaReturnType :: [Type]
lambdaReturnType =
              Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_p)
                [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType [Param Type]
forward_params,
            lambdaBody :: Body SOACS
lambdaBody =
              Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody
                (Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_p)))
                ( Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_p))
                    [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [VName] -> [SubExpRes]
varsRes ((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
forward_params)
                )
          }

  let lam2 =
        Lambda
          { lambdaParams :: [LParam SOACS]
lambdaParams =
              Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_p)
                [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Bool] -> [Param Type] -> [Param Type]
forall a. [Bool] -> [a] -> [a]
pick [Bool]
inp_c_real_map (Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_c)),
            lambdaReturnType :: [Type]
lambdaReturnType =
              Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_c)
                [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_p),
            lambdaBody :: Body SOACS
lambdaBody =
              Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody
                ( Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_p))
                    Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Lambda SOACS -> [VName] -> [VName] -> Lambda SOACS -> Stms SOACS
forall rep a.
(Buildable rep, Ord a) =>
Lambda rep -> [a] -> [a] -> Lambda rep -> Stms rep
composeBinds
                      (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_p)
                      [VName]
out_post_p
                      (Input -> VName
SOAC.inputArray (Input -> VName) -> [Input] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Input]
inp_c)
                      (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_c)
                    Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_c))
                )
                ( Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_c))
                    [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_p))
                )
          }

  post_forward_params <- forM
    ( zip
        (bodyResult (lambdaBody (scremaPostLambda form_p)))
        (lambdaReturnType (scremaPostLambda form_p))
    )
    $ \(SubExpRes
res, Type
t) ->
      Name -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
Name -> dec -> m (Param dec)
newParam (Name -> (VName -> Name) -> Maybe VName -> Name
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String -> Name
nameFromString String
"x") VName -> Name
baseName (SubExpRes -> Maybe VName
subExpResVName SubExpRes
res)) Type
t

  let lam3 =
        Lambda
          { lambdaParams :: [LParam SOACS]
lambdaParams =
              Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_c) [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type]
post_forward_params,
            lambdaReturnType :: [Type]
lambdaReturnType =
              Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_c)
                [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall dec. Typed dec => Param dec -> Type
paramType [Param Type]
post_forward_params,
            lambdaBody :: Body SOACS
lambdaBody =
              Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody
                (Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_c)))
                ( Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody (ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaPostLambda ScremaForm SOACS
form_c))
                    [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [VName] -> [SubExpRes]
varsRes ((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
post_forward_params)
                )
          }
  -- Deduplicate inp_r: when inp_p and inp_c_real share inputs, the
  -- combined input list contains duplicates.  Remove them by keeping
  -- only one copy of each unique input and replacing the extra lambda
  -- parameters in lam1 with let-bindings.
  let nil_post = [LParam SOACS] -> [t] -> Body SOACS -> GLambda SOACS t
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [] [] (Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
forall a. Monoid a => a
mempty [])
      (inp_r', form_lam1') = dedupInput inp_r (ScremaForm lam1 [] [] nil_post)
      lam1' = ScremaForm SOACS -> Lambda SOACS
forall rep. ScremaForm rep -> Lambda rep
scremaLambda ScremaForm SOACS
form_lam1'
  pure
    ( SuperScrema
        w
        inp_r'
        lam1'
        (scremaScans form_p)
        (scremaReduces form_p)
        lam2
        (scremaScans form_c)
        (scremaReduces form_c)
        lam3,
      out_red_p <> out_red_c <> out_post_c <> out_post_p
    )
  where
    inputFromOutput :: Input -> Bool
inputFromOutput Input
inp = Input -> VName
SOAC.inputArray Input
inp VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
out_p

-- | Moves the last scan and reduce from the super screma to the top
-- of the super screma.
moveRedScanSuperScrema ::
  (MonadFail m, MonadFreshNames m) =>
  SuperScrema SOACS ->
  m (SuperScrema SOACS)
moveRedScanSuperScrema :: forall (m :: * -> *).
(MonadFail m, MonadFreshNames m) =>
SuperScrema SOACS -> m (SuperScrema SOACS)
moveRedScanSuperScrema super_screma :: SuperScrema SOACS
super_screma@(SuperScrema SubExp
_ [Input]
_ Lambda SOACS
_ [Scan SOACS]
_ [Reduce SOACS]
_ Lambda SOACS
_ [] [] Lambda SOACS
_) =
  SuperScrema SOACS -> m (SuperScrema SOACS)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SuperScrema SOACS
super_screma
moveRedScanSuperScrema SuperScrema SOACS
super_screma = do
  ((scan_red_inp_c, scan_red_lam', _), (_, map_lam', _)) <-
    Int
-> [Int]
-> Lambda SOACS
-> [()]
-> m (([Int], Lambda SOACS, [()]), ([Int], Lambda SOACS, [()]))
forall (m :: * -> *) inp out.
MonadFail m =>
Int
-> [inp]
-> Lambda SOACS
-> [out]
-> m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
splitAtLambdaByRes
      ([Scan SOACS] -> Int
forall rep. [Scan rep] -> Int
scanResults [Scan SOACS]
scan' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
red')
      [Int]
inp_c
      Lambda SOACS
lam'
      (Int -> () -> [()]
forall a. Int -> a -> [a]
replicate ([Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Type] -> Int) -> [Type] -> Int
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam') ())
  renamed_scan_red_lam' <- renameLambda scan_red_lam'
  let (scan_res', red_res') =
        splitAt (scanResults scan')
          . bodyResult
          $ lambdaBody renamed_scan_red_lam'
      (scan_ts', red_ts') =
        splitAt (scanResults scan') $
          lambdaReturnType renamed_scan_red_lam'
      binds = Lambda SOACS -> [Int] -> [Int] -> Lambda SOACS -> Stms SOACS
forall rep a.
(Buildable rep, Ord a) =>
Lambda rep -> [a] -> [a] -> Lambda rep -> Stms rep
composeBinds Lambda SOACS
lam [Int]
out_p [Int]
scan_red_inp_c Lambda SOACS
renamed_scan_red_lam'
      stms' = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
renamed_scan_red_lam'
      new_scan = [Scan SOACS]
scan [Scan SOACS] -> [Scan SOACS] -> [Scan SOACS]
forall a. Semigroup a => a -> a -> a
<> [Scan SOACS]
scan'
      new_red = [Reduce SOACS]
red [Reduce SOACS] -> [Reduce SOACS] -> [Reduce SOACS]
forall a. Semigroup a => a -> a -> a
<> [Reduce SOACS]
red'
      new_ts = [Type]
scan_ts [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> [Type]
scan_ts' [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> [Type]
red_ts [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> [Type]
red_ts' [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> [Type]
map_ts
      new_pars = Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam
      new_res = [SubExpRes]
scan_res [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [SubExpRes]
scan_res' [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [SubExpRes]
red_res [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [SubExpRes]
red_res' [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [SubExpRes]
map_res
      new_body = Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody (Stms SOACS
stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stms SOACS
binds Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stms SOACS
stms') [SubExpRes]
new_res
      new_lam = [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
new_pars [Type]
new_ts Body SOACS
new_body
      (scan_pars', map_pars') =
        splitAt (scanResults scan) (lambdaParams lam')

  extra_scan_pars' <- mapM (newParam "x") scan_ts'

  let new_pars' = [Param Type]
scan_pars' [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type]
extra_scan_pars' [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type]
map_pars'
      new_ts' = [Type]
scan_ts' [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
map_lam'
      new_stms' = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
map_lam'
      new_res' =
        [VName] -> [SubExpRes]
varsRes ((Param Type -> VName) -> [Param Type] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> VName
forall dec. Param dec -> VName
paramName [Param Type]
extra_scan_pars')
          [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
map_lam')
      new_body' = Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
new_stms' [SubExpRes]
new_res'
      new_lam' = [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
new_pars' [Type]
new_ts' Body SOACS
new_body'

  pure $
    SuperScrema w inp new_lam new_scan new_red new_lam' [] [] lam''
  where
    stms :: Stms SOACS
stms = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam
    out_p :: [Int]
out_p = [Int
0 .. [SubExpRes] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    ([Int]
scan_out, [Int]
_, [Int]
map_out) =
      Int -> Int -> [Int] -> ([Int], [Int], [Int])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3
        ([Scan SOACS] -> Int
forall rep. [Scan rep] -> Int
scanResults [Scan SOACS]
scan)
        ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
red)
        [Int
0 .. [SubExpRes] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    ([SubExpRes]
scan_res, [SubExpRes]
red_res, [SubExpRes]
map_res) =
      Int
-> Int -> [SubExpRes] -> ([SubExpRes], [SubExpRes], [SubExpRes])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3
        ([Scan SOACS] -> Int
forall rep. [Scan rep] -> Int
scanResults [Scan SOACS]
scan)
        ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
red)
        (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam)
    ([Type]
scan_ts, [Type]
red_ts, [Type]
map_ts) =
      Int -> Int -> [Type] -> ([Type], [Type], [Type])
forall a. Int -> Int -> [a] -> ([a], [a], [a])
splitAt3
        ([Scan SOACS] -> Int
forall rep. [Scan rep] -> Int
scanResults [Scan SOACS]
scan)
        ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
red)
        (Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam)
    inp_c :: [Int]
inp_c = [Int]
scan_out [Int] -> [Int] -> [Int]
forall a. Semigroup a => a -> a -> a
<> [Int]
map_out

    SuperScrema SubExp
w [Input]
inp Lambda SOACS
lam [Scan SOACS]
scan [Reduce SOACS]
red Lambda SOACS
lam' [Scan SOACS]
scan' [Reduce SOACS]
red' Lambda SOACS
lam'' =
      SuperScrema SOACS
super_screma

-- | Moves all work done in the last lambda to middle lambda of the
-- super screma.
moveLastSuperScrema ::
  (MonadFreshNames m) =>
  SuperScrema SOACS ->
  m (SuperScrema SOACS)
moveLastSuperScrema :: forall (m :: * -> *).
MonadFreshNames m =>
SuperScrema SOACS -> m (SuperScrema SOACS)
moveLastSuperScrema (SuperScrema SubExp
w [Input]
inp Lambda SOACS
lam [Scan SOACS]
scan [Reduce SOACS]
red Lambda SOACS
lam' [] [] Lambda SOACS
lam'') = do
  temp_lam'' <- Lambda SOACS -> m (Lambda SOACS)
forall rep (m :: * -> *) t.
(Renameable rep, MonadFreshNames m, Rename t) =>
GLambda rep t -> m (GLambda rep t)
renameLambda Lambda SOACS
lam''
  let new_pars = Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam'
      new_ts = Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam''
      out_p = [Int
0 .. [SubExpRes] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam') Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      inp_c = [Int]
out_p
      binds = Lambda SOACS -> [Int] -> [Int] -> Lambda SOACS -> Stms SOACS
forall rep a.
(Buildable rep, Ord a) =>
Lambda rep -> [a] -> [a] -> Lambda rep -> Stms rep
composeBinds Lambda SOACS
lam' [Int]
out_p [Int]
inp_c Lambda SOACS
temp_lam''
      stms' = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam'
      stms'' = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
temp_lam''
      new_stms = Stms SOACS
stms' Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stms SOACS
binds Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stms SOACS
stms''
      new_res = Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
temp_lam''
      new_body = Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
new_stms [SubExpRes]
new_res
      new_lam' = [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
new_pars [Type]
new_ts Body SOACS
new_body

  new_lam'' <- mkIdentityLambda $ lambdaReturnType lam''
  pure $
    SuperScrema w inp lam scan red new_lam' [] [] new_lam''
moveLastSuperScrema SuperScrema SOACS
_ =
  String -> m (SuperScrema SOACS)
forall a. HasCallStack => String -> a
error String
"moveLastSuperScrema must not have any scans or reduce operation in the end."

-- | Find all Accumulator parameters.
parAccs :: Lambda SOACS -> [Type]
parAccs :: Lambda SOACS -> [Type]
parAccs = (Type -> Bool) -> [Type] -> [Type]
forall a. (a -> Bool) -> [a] -> [a]
filter Type -> Bool
forall shape o. TypeBase shape o -> Bool
isAcc ([Type] -> [Type])
-> (Lambda SOACS -> [Type]) -> Lambda SOACS -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param Type -> Type) -> [Param Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Type
forall t. Typed t => t -> Type
typeOf ([Param Type] -> [Type])
-> (Lambda SOACS -> [Param Type]) -> Lambda SOACS -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda SOACS -> [Param Type]
Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams

-- | Check if the lambda parameters have overlapping accumulators.
parAccsOverlap :: Lambda SOACS -> Lambda SOACS -> Bool
parAccsOverlap :: Lambda SOACS -> Lambda SOACS -> Bool
parAccsOverlap Lambda SOACS
lam = (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Type -> [Type] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Type]
accs) ([Type] -> Bool)
-> (Lambda SOACS -> [Type]) -> Lambda SOACS -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lambda SOACS -> [Type]
parAccs
  where
    accs :: [Type]
accs = Lambda SOACS -> [Type]
parAccs Lambda SOACS
lam

-- | Check if the lambdas have parameters that overlap due to
-- consumption.
consumedOverlap :: Lambda SOACS -> Lambda SOACS -> Bool
consumedOverlap :: Lambda SOACS -> Lambda SOACS -> Bool
consumedOverlap =
  (Names -> Names -> Bool)
-> (Lambda SOACS -> Names) -> Lambda SOACS -> Lambda SOACS -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
on Names -> Names -> Bool
namesIntersect (Lambda (Aliases SOACS) -> Names
forall rep. Aliased rep => Lambda rep -> Names
consumedByLambda (Lambda (Aliases SOACS) -> Names)
-> (Lambda SOACS -> Lambda (Aliases SOACS))
-> Lambda SOACS
-> Names
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AliasTable -> Lambda SOACS -> Lambda (Aliases SOACS)
forall rep t.
AliasableRep rep =>
AliasTable -> GLambda rep t -> GLambda (Aliases rep) t
analyseLambda AliasTable
forall a. Monoid a => a
mempty)

-- | Split a lambda at some index for the result and construct two
-- lambda which do not have overlapping results but may have
-- overlapping parameters.
splitAtLambdaByRes ::
  (MonadFail m) =>
  Int ->
  [inp] ->
  Lambda SOACS ->
  [out] ->
  m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
splitAtLambdaByRes :: forall (m :: * -> *) inp out.
MonadFail m =>
Int
-> [inp]
-> Lambda SOACS
-> [out]
-> m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
splitAtLambdaByRes Int
i [inp]
inp Lambda SOACS
lam [out]
out = do
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when
    (Lambda SOACS -> Lambda SOACS -> Bool
consumedOverlap Lambda SOACS
new_lam Lambda SOACS
new_lam')
    (String -> m ()
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Can not fuse due to overlap consumption in pre and post.")
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when
    (Lambda SOACS -> Lambda SOACS -> Bool
parAccsOverlap Lambda SOACS
new_lam Lambda SOACS
new_lam')
    (String -> m ()
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Can not fuse due to overlap in parameter accumalators.")
  (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
-> m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out]))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (([inp]
new_inp, Lambda SOACS
new_lam, [out]
new_out), ([inp]
new_inp', Lambda SOACS
new_lam', [out]
new_out'))
  where
    new_lam :: Lambda SOACS
new_lam = [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
new_pars [Type]
new_ts Body SOACS
new_body
    new_lam' :: Lambda SOACS
new_lam' = [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [Param Type]
[LParam SOACS]
new_pars' [Type]
new_ts' Body SOACS
new_body'
    pars :: [LParam SOACS]
pars = Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam
    stms :: Stms SOACS
stms = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam
    (Names
new_deps, Stms SOACS
new_stms) = Names -> Stms SOACS -> (Names, Stms SOACS)
forall rep. Buildable rep => Names -> Stms rep -> (Names, Stms rep)
eliminateWithDeps ([SubExpRes] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExpRes]
new_res) Stms SOACS
stms
    (Names
new_deps', Stms SOACS
new_stms') = Names -> Stms SOACS -> (Names, Stms SOACS)
forall rep. Buildable rep => Names -> Stms rep -> (Names, Stms rep)
eliminateWithDeps ([SubExpRes] -> Names
forall a. FreeIn a => a -> Names
freeIn [SubExpRes]
new_res') Stms SOACS
stms
    new_body :: Body SOACS
new_body = Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
new_stms [SubExpRes]
new_res
    new_body' :: Body SOACS
new_body' = Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody Stms SOACS
new_stms' [SubExpRes]
new_res'
    inBody :: Names -> (Param dec, b) -> Bool
inBody Names
deps = (VName -> Names -> Bool
`nameIn` Names
deps) (VName -> Bool)
-> ((Param dec, b) -> VName) -> (Param dec, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param dec -> VName
forall dec. Param dec -> VName
paramName (Param dec -> VName)
-> ((Param dec, b) -> Param dec) -> (Param dec, b) -> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param dec, b) -> Param dec
forall a b. (a, b) -> a
fst
    removePars :: Names -> ([Param Type], [inp])
removePars Names
deps = [(Param Type, inp)] -> ([Param Type], [inp])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Param Type, inp)] -> ([Param Type], [inp]))
-> [(Param Type, inp)] -> ([Param Type], [inp])
forall a b. (a -> b) -> a -> b
$ ((Param Type, inp) -> Bool)
-> [(Param Type, inp)] -> [(Param Type, inp)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Names -> (Param Type, inp) -> Bool
forall {dec} {b}. Names -> (Param dec, b) -> Bool
inBody Names
deps) ([(Param Type, inp)] -> [(Param Type, inp)])
-> [(Param Type, inp)] -> [(Param Type, inp)]
forall a b. (a -> b) -> a -> b
$ [Param Type] -> [inp] -> [(Param Type, inp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Param Type]
pars [inp]
inp
    ([Param Type]
new_pars, [inp]
new_inp) = Names -> ([Param Type], [inp])
removePars Names
new_deps
    ([Param Type]
new_pars', [inp]
new_inp') = Names -> ([Param Type], [inp])
removePars Names
new_deps'
    ([SubExpRes]
new_res, [SubExpRes]
new_res') = Int -> [SubExpRes] -> ([SubExpRes], [SubExpRes])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
i ([SubExpRes] -> ([SubExpRes], [SubExpRes]))
-> [SubExpRes] -> ([SubExpRes], [SubExpRes])
forall a b. (a -> b) -> a -> b
$ Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam
    ([Type]
new_ts, [Type]
new_ts') = Int -> [Type] -> ([Type], [Type])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
i ([Type] -> ([Type], [Type])) -> [Type] -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
lam
    ([out]
new_out, [out]
new_out') = Int -> [out] -> ([out], [out])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
i [out]
out

-- | Create a mapping from lambda results expression to their output
-- array.
resToOut :: [VName] -> Lambda SOACS -> SubExpRes -> VName
resToOut :: [VName] -> Lambda SOACS -> SubExpRes -> VName
resToOut [VName]
out Lambda SOACS
lam = (Map SubExpRes VName
m Map SubExpRes VName -> SubExpRes -> VName
forall k a. Ord k => Map k a -> k -> a
M.!)
  where
    m :: Map SubExpRes VName
m = [(SubExpRes, VName)] -> Map SubExpRes VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(SubExpRes, VName)] -> Map SubExpRes VName)
-> [(SubExpRes, VName)] -> Map SubExpRes VName
forall a b. (a -> b) -> a -> b
$ ([SubExpRes] -> [VName] -> [(SubExpRes, VName)])
-> [VName] -> [SubExpRes] -> [(SubExpRes, VName)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [SubExpRes] -> [VName] -> [(SubExpRes, VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
out ([SubExpRes] -> [(SubExpRes, VName)])
-> [SubExpRes] -> [(SubExpRes, VName)]
forall a b. (a -> b) -> a -> b
$ Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
lam

-- | Create a mapping from lambda parameter names to their input
-- array.
parToInp :: [SOAC.Input] -> Lambda SOACS -> VName -> Maybe SOAC.Input
parToInp :: [Input] -> Lambda SOACS -> VName -> Maybe Input
parToInp [Input]
inp Lambda SOACS
lam = (VName -> Map VName Input -> Maybe Input)
-> Map VName Input -> VName -> Maybe Input
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> Map VName Input -> Maybe Input
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Map VName Input
m
  where
    m :: Map VName Input
m = [(VName, Input)] -> Map VName Input
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Input)] -> Map VName Input)
-> [(VName, Input)] -> Map VName Input
forall a b. (a -> b) -> a -> b
$ ([VName] -> [Input] -> [(VName, Input)])
-> [Input] -> [VName] -> [(VName, Input)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [VName] -> [Input] -> [(VName, Input)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Input]
inp ([VName] -> [(VName, Input)]) -> [VName] -> [(VName, Input)]
forall a b. (a -> b) -> a -> b
$ Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> VName) -> [Param Type] -> [VName]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
lam

-- | Turn a SuperScrema into Screma.
toScrema ::
  SuperScrema SOACS ->
  ([SOAC.Input], ScremaForm SOACS)
toScrema :: SuperScrema SOACS -> ([Input], ScremaForm SOACS)
toScrema (SuperScrema SubExp
_ [Input]
inp Lambda SOACS
lam [Scan SOACS]
scan [Reduce SOACS]
red Lambda SOACS
lam' [Scan SOACS]
_ [Reduce SOACS]
_ Lambda SOACS
_) =
  ([Input]
inp, Lambda SOACS
-> [Scan SOACS]
-> [Reduce SOACS]
-> Lambda SOACS
-> ScremaForm SOACS
forall rep.
Lambda rep
-> [Scan rep] -> [Reduce rep] -> Lambda rep -> ScremaForm rep
ScremaForm Lambda SOACS
lam [Scan SOACS]
scan [Reduce SOACS]
red Lambda SOACS
lam')

-- | Try to force the post lambda to be an indentity lambda.
tryIdentityPost ::
  (MonadFreshNames m) => ScremaForm SOACS -> m (ScremaForm SOACS)
tryIdentityPost :: forall (m :: * -> *).
MonadFreshNames m =>
ScremaForm SOACS -> m (ScremaForm SOACS)
tryIdentityPost (ScremaForm Lambda SOACS
pre_lam [] [Reduce SOACS]
reds Lambda SOACS
post_lam)
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Bool
forall rep. Lambda rep -> Bool
isIdentityLambda Lambda SOACS
post_lam = do
      new_post_lam <- [Type] -> m (Lambda SOACS)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m) =>
[Type] -> m (Lambda rep)
mkIdentityLambda ([Type] -> m (Lambda SOACS)) -> [Type] -> m (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
post_lam
      pure (ScremaForm new_pre_lam [] reds new_post_lam)
  where
    out_p :: [Int]
out_p = [Int
0 .. [SubExpRes] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
pre_lam) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    inp_c :: [Int]
inp_c = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
reds) [Int]
out_p
    binds :: Stms SOACS
binds = Lambda SOACS -> [Int] -> [Int] -> Lambda SOACS -> Stms SOACS
forall rep a.
(Buildable rep, Ord a) =>
Lambda rep -> [a] -> [a] -> Lambda rep -> Stms rep
composeBinds Lambda SOACS
pre_lam [Int]
out_p [Int]
inp_c Lambda SOACS
post_lam
    pre_stms :: Stms SOACS
pre_stms = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
pre_lam
    post_stms :: Stms SOACS
post_stms = Body SOACS -> Stms SOACS
forall rep res. GBody rep res -> Stms rep
bodyStms (Body SOACS -> Stms SOACS) -> Body SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
post_lam
    post_res :: [SubExpRes]
post_res = Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
post_lam
    pre_res :: [SubExpRes]
pre_res = Body SOACS -> [SubExpRes]
forall rep res. GBody rep res -> [res]
bodyResult (Body SOACS -> [SubExpRes]) -> Body SOACS -> [SubExpRes]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep t. GLambda rep t -> Body rep
lambdaBody Lambda SOACS
pre_lam
    post_ts :: [Type]
post_ts = Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
post_lam
    pre_ts :: [Type]
pre_ts = Lambda SOACS -> [Type]
forall rep t. GLambda rep t -> [t]
lambdaReturnType Lambda SOACS
pre_lam
    res :: [SubExpRes]
res = Int -> [SubExpRes] -> [SubExpRes]
forall a. Int -> [a] -> [a]
take ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
reds) [SubExpRes]
pre_res [SubExpRes] -> [SubExpRes] -> [SubExpRes]
forall a. Semigroup a => a -> a -> a
<> [SubExpRes]
post_res
    ts :: [Type]
ts = Int -> [Type] -> [Type]
forall a. Int -> [a] -> [a]
take ([Reduce SOACS] -> Int
forall rep. [Reduce rep] -> Int
redResults [Reduce SOACS]
reds) [Type]
pre_ts [Type] -> [Type] -> [Type]
forall a. Semigroup a => a -> a -> a
<> [Type]
post_ts
    new_pre_lam :: Lambda SOACS
new_pre_lam =
      Lambda
        { lambdaParams :: [LParam SOACS]
lambdaParams = Lambda SOACS -> [LParam SOACS]
forall rep t. GLambda rep t -> [LParam rep]
lambdaParams Lambda SOACS
pre_lam,
          lambdaReturnType :: [Type]
lambdaReturnType = [Type]
ts,
          lambdaBody :: Body SOACS
lambdaBody = Stms SOACS -> [SubExpRes] -> Body SOACS
forall res. IsResult res => Stms SOACS -> [res] -> GBody SOACS res
forall rep res.
(Buildable rep, IsResult res) =>
Stms rep -> [res] -> GBody rep res
mkBody (Stms SOACS
pre_stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stms SOACS
binds Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stms SOACS
post_stms) [SubExpRes]
res
        }
tryIdentityPost ScremaForm SOACS
form = ScremaForm SOACS -> m (ScremaForm SOACS)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScremaForm SOACS
form

-- Tries to fuse two Scremas into one.
fuseScrema ::
  (MonadFail m, MonadFreshNames m, HasScope SOACS m) =>
  SubExp ->
  [SOAC.Input] ->
  ScremaForm SOACS ->
  [VName] ->
  [SOAC.Input] ->
  ScremaForm SOACS ->
  [VName] ->
  m ([SOAC.Input], ScremaForm SOACS, [VName])
fuseScrema :: forall (m :: * -> *).
(MonadFail m, MonadFreshNames m, HasScope SOACS m) =>
SubExp
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> m ([Input], ScremaForm SOACS, [VName])
fuseScrema SubExp
w [Input]
inp_p ScremaForm SOACS
form_p [VName]
out_p [Input]
inp_c ScremaForm SOACS
form_c [VName]
out_c = do
  ScremaForm SOACS -> [VName] -> [Input] -> ScremaForm SOACS -> m ()
forall (m :: * -> *).
MonadFail m =>
ScremaForm SOACS -> [VName] -> [Input] -> ScremaForm SOACS -> m ()
fusible ScremaForm SOACS
form_p [VName]
out_p [Input]
inp_c ScremaForm SOACS
form_c
  (super_screma, new_out) <- SubExp
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> m (SuperScrema SOACS, [VName])
forall (m :: * -> *).
MonadFreshNames m =>
SubExp
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> [Input]
-> ScremaForm SOACS
-> [VName]
-> m (SuperScrema SOACS, [VName])
fuseSuperScrema SubExp
w [Input]
inp_p ScremaForm SOACS
form_p [VName]
out_p [Input]
inp_c ScremaForm SOACS
form_c [VName]
out_c
  (new_inp, form') <-
    fmap (second prunePreLambdaResults . toScrema) $
      moveRedScanSuperScrema super_screma
        >>= moveLastSuperScrema
  form <- tryIdentityPost form'
  pure (new_inp, form, new_out)