{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}

-- | It is well known that fully parallel loops can always be
-- interchanged inwards with a sequential loop.  This module
-- implements that transformation.
--
-- This is also where we implement loop-switching (for branches),
-- which is semantically similar to interchange.
module Futhark.Pass.ExtractKernels.Interchange
  ( SeqLoop (..),
    interchangeLoops,
    Branch (..),
    interchangeBranch,
    WithAccStm (..),
    interchangeWithAcc,
  )
where

import Control.Monad.Identity
import Data.List (find)
import Data.Maybe
import Futhark.IR.SOACS
import Futhark.MonadFreshNames
import Futhark.Pass.ExtractKernels.Distribution
  ( KernelNest,
    LoopNesting (..),
    kernelNestLoops,
    scopeOfKernelNest,
  )
import Futhark.Tools
import Futhark.Transform.Rename
import Futhark.Util (splitFromEnd)

-- | An encoding of a sequential do-loop with no existential context,
-- alongside its result pattern.
data SeqLoop
  = SeqLoop [Int] (Pat Type) [(FParam SOACS, SubExp)] (LoopForm SOACS) (Body SOACS)

loopPerm :: SeqLoop -> [Int]
loopPerm :: SeqLoop -> [Int]
loopPerm (SeqLoop [Int]
perm Pat Type
_ [(FParam SOACS, SubExp)]
_ LoopForm SOACS
_ Body SOACS
_) = [Int]
perm

seqLoopStm :: SeqLoop -> Stm SOACS
seqLoopStm :: SeqLoop -> Stm SOACS
seqLoopStm (SeqLoop [Int]
_ Pat Type
pat [(FParam SOACS, SubExp)]
merge LoopForm SOACS
form Body SOACS
body) =
  Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat Type
Pat (LetDec SOACS)
pat (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$ [(FParam SOACS, SubExp)]
-> LoopForm SOACS -> Body SOACS -> Exp SOACS
forall rep.
[(FParam rep, SubExp)] -> LoopForm rep -> Body rep -> Exp rep
DoLoop [(FParam SOACS, SubExp)]
merge LoopForm SOACS
form Body SOACS
body

interchangeLoop ::
  (MonadBuilder m, LocalScope SOACS m) =>
  (VName -> Maybe VName) ->
  SeqLoop ->
  LoopNesting ->
  m SeqLoop
interchangeLoop :: (VName -> Maybe VName) -> SeqLoop -> LoopNesting -> m SeqLoop
interchangeLoop
  VName -> Maybe VName
isMapParameter
  (SeqLoop [Int]
perm Pat Type
loop_pat [(FParam SOACS, SubExp)]
merge LoopForm SOACS
form Body SOACS
body)
  (MapNesting Pat Type
pat StmAux ()
aux SubExp
w [(Param Type, VName)]
params_and_arrs) = do
    [(Param (TypeBase Shape Uniqueness), SubExp)]
merge_expanded <-
      Scope SOACS
-> m [(Param (TypeBase Shape Uniqueness), SubExp)]
-> m [(Param (TypeBase Shape Uniqueness), SubExp)]
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams ([Param Type] -> Scope SOACS) -> [Param Type] -> Scope SOACS
forall a b. (a -> b) -> a -> b
$ ((Param Type, VName) -> Param Type)
-> [(Param Type, VName)] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map (Param Type, VName) -> Param Type
forall a b. (a, b) -> a
fst [(Param Type, VName)]
params_and_arrs) (m [(Param (TypeBase Shape Uniqueness), SubExp)]
 -> m [(Param (TypeBase Shape Uniqueness), SubExp)])
-> m [(Param (TypeBase Shape Uniqueness), SubExp)]
-> m [(Param (TypeBase Shape Uniqueness), SubExp)]
forall a b. (a -> b) -> a -> b
$
        ((Param (TypeBase Shape Uniqueness), SubExp)
 -> m (Param (TypeBase Shape Uniqueness), SubExp))
-> [(Param (TypeBase Shape Uniqueness), SubExp)]
-> m [(Param (TypeBase Shape Uniqueness), SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Param (TypeBase Shape Uniqueness), SubExp)
-> m (Param (TypeBase Shape Uniqueness), SubExp)
expand [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam SOACS, SubExp)]
merge

    let loop_pat_expanded :: Pat Type
loop_pat_expanded =
          [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ (PatElem Type -> PatElem Type) -> [PatElem Type] -> [PatElem Type]
forall a b. (a -> b) -> [a] -> [b]
map PatElem Type -> PatElem Type
expandPatElem ([PatElem Type] -> [PatElem Type])
-> [PatElem Type] -> [PatElem Type]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
loop_pat
        new_params :: [Param Type]
new_params =
          [Attrs -> VName -> Type -> Param Type
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
attrs VName
pname (Type -> Param Type) -> Type -> Param Type
forall a b. (a -> b) -> a -> b
$ TypeBase Shape Uniqueness -> Type
forall shape.
TypeBase shape Uniqueness -> TypeBase shape NoUniqueness
fromDecl TypeBase Shape Uniqueness
ptype | (Param Attrs
attrs VName
pname TypeBase Shape Uniqueness
ptype, SubExp
_) <- [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam SOACS, SubExp)]
merge]
        new_arrs :: [VName]
new_arrs = ((Param (TypeBase Shape Uniqueness), SubExp) -> VName)
-> [(Param (TypeBase Shape Uniqueness), SubExp)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Param (TypeBase Shape Uniqueness) -> VName
forall dec. Param dec -> VName
paramName (Param (TypeBase Shape Uniqueness) -> VName)
-> ((Param (TypeBase Shape Uniqueness), SubExp)
    -> Param (TypeBase Shape Uniqueness))
-> (Param (TypeBase Shape Uniqueness), SubExp)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param (TypeBase Shape Uniqueness), SubExp)
-> Param (TypeBase Shape Uniqueness)
forall a b. (a, b) -> a
fst) [(Param (TypeBase Shape Uniqueness), SubExp)]
merge_expanded
        rettype :: [Type]
rettype = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall u. TypeBase Shape u -> TypeBase Shape u
rowType ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat Type
loop_pat_expanded

    -- If the map consumes something that is bound outside the loop
    -- (i.e. is not a merge parameter), we have to copy() it.  As a
    -- small simplification, we just remove the parameter outright if
    -- it is not used anymore.  This might happen if the parameter was
    -- used just as the inital value of a merge parameter.
    (([Param Type]
params', [VName]
arrs'), Stms SOACS
pre_copy_stms) <-
      Builder SOACS ([Param Type], [VName])
-> m (([Param Type], [VName]), Stms SOACS)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder SOACS ([Param Type], [VName])
 -> m (([Param Type], [VName]), Stms SOACS))
-> Builder SOACS ([Param Type], [VName])
-> m (([Param Type], [VName]), Stms SOACS)
forall a b. (a -> b) -> a -> b
$
        Scope SOACS
-> Builder SOACS ([Param Type], [VName])
-> Builder SOACS ([Param Type], [VName])
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param Type]
new_params) (Builder SOACS ([Param Type], [VName])
 -> Builder SOACS ([Param Type], [VName]))
-> Builder SOACS ([Param Type], [VName])
-> Builder SOACS ([Param Type], [VName])
forall a b. (a -> b) -> a -> b
$
          [(Param Type, VName)] -> ([Param Type], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Param Type, VName)] -> ([Param Type], [VName]))
-> ([Maybe (Param Type, VName)] -> [(Param Type, VName)])
-> [Maybe (Param Type, VName)]
-> ([Param Type], [VName])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (Param Type, VName)] -> [(Param Type, VName)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Param Type, VName)] -> ([Param Type], [VName]))
-> BuilderT SOACS (State VNameSource) [Maybe (Param Type, VName)]
-> Builder SOACS ([Param Type], [VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Param Type, VName)
 -> BuilderT SOACS (State VNameSource) (Maybe (Param Type, VName)))
-> [(Param Type, VName)]
-> BuilderT SOACS (State VNameSource) [Maybe (Param Type, VName)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Param Type, VName)
-> BuilderT SOACS (State VNameSource) (Maybe (Param Type, VName))
copyOrRemoveParam [(Param Type, VName)]
params_and_arrs

    let lam :: Lambda SOACS
lam = [LParam SOACS] -> Body SOACS -> [Type] -> Lambda SOACS
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda ([Param Type]
params' [Param Type] -> [Param Type] -> [Param Type]
forall a. Semigroup a => a -> a -> a
<> [Param Type]
new_params) Body SOACS
body [Type]
rettype
        map_stm :: Stm SOACS
map_stm =
          Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat Type
Pat (LetDec SOACS)
loop_pat_expanded StmAux ()
StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
            Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w ([VName]
arrs' [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
new_arrs) (Lambda SOACS -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda SOACS
lam)
        res :: Result
res = [VName] -> Result
varsRes ([VName] -> Result) -> [VName] -> Result
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
loop_pat_expanded
        pat' :: Pat Type
pat' = [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ [Int] -> [PatElem Type] -> [PatElem Type]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm ([PatElem Type] -> [PatElem Type])
-> [PatElem Type] -> [PatElem Type]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat

    SeqLoop -> m SeqLoop
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SeqLoop -> m SeqLoop) -> SeqLoop -> m SeqLoop
forall a b. (a -> b) -> a -> b
$
      [Int]
-> Pat Type
-> [(FParam SOACS, SubExp)]
-> LoopForm SOACS
-> Body SOACS
-> SeqLoop
SeqLoop [Int]
perm Pat Type
pat' [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam SOACS, SubExp)]
merge_expanded LoopForm SOACS
form (Body SOACS -> SeqLoop) -> Body SOACS -> SeqLoop
forall a b. (a -> b) -> a -> b
$
        Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody (Stms SOACS
pre_copy_stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
map_stm) Result
res
    where
      free_in_body :: Names
free_in_body = Body SOACS -> Names
forall a. FreeIn a => a -> Names
freeIn Body SOACS
body

      copyOrRemoveParam :: (Param Type, VName)
-> BuilderT SOACS (State VNameSource) (Maybe (Param Type, VName))
copyOrRemoveParam (Param Type
param, VName
arr)
        | Bool -> Bool
not (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
param VName -> Names -> Bool
`nameIn` Names
free_in_body) =
            Maybe (Param Type, VName)
-> BuilderT SOACS (State VNameSource) (Maybe (Param Type, VName))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Param Type, VName)
forall a. Maybe a
Nothing
        | Bool
otherwise =
            Maybe (Param Type, VName)
-> BuilderT SOACS (State VNameSource) (Maybe (Param Type, VName))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Param Type, VName)
 -> BuilderT SOACS (State VNameSource) (Maybe (Param Type, VName)))
-> Maybe (Param Type, VName)
-> BuilderT SOACS (State VNameSource) (Maybe (Param Type, VName))
forall a b. (a -> b) -> a -> b
$ (Param Type, VName) -> Maybe (Param Type, VName)
forall a. a -> Maybe a
Just (Param Type
param, VName
arr)

      expandedInit :: [Char] -> SubExp -> m SubExp
expandedInit [Char]
_ (Var VName
v)
        | Just VName
arr <- VName -> Maybe VName
isMapParameter VName
v =
            SubExp -> m SubExp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SubExp -> m SubExp) -> SubExp -> m SubExp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
arr
      expandedInit [Char]
param_name SubExp
se =
        [Char] -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp ([Char]
param_name [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_expanded_init") (Exp (Rep m) -> m SubExp) -> Exp (Rep m) -> m SubExp
forall a b. (a -> b) -> a -> b
$
          BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ Shape -> SubExp -> BasicOp
Replicate ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) SubExp
se

      expand :: (Param (TypeBase Shape Uniqueness), SubExp)
-> m (Param (TypeBase Shape Uniqueness), SubExp)
expand (Param (TypeBase Shape Uniqueness)
merge_param, SubExp
merge_init) = do
        Param (TypeBase Shape Uniqueness)
expanded_param <-
          [Char]
-> TypeBase Shape Uniqueness
-> m (Param (TypeBase Shape Uniqueness))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam ([Char]
param_name [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_expanded") (TypeBase Shape Uniqueness
 -> m (Param (TypeBase Shape Uniqueness)))
-> TypeBase Shape Uniqueness
-> m (Param (TypeBase Shape Uniqueness))
forall a b. (a -> b) -> a -> b
$
            -- FIXME: Unique here is a hack to make sure the copy from
            -- makeCopyInitial is not prematurely simplified away.
            -- It'd be better to fix this somewhere else...
            TypeBase Shape Uniqueness
-> Shape -> Uniqueness -> TypeBase Shape Uniqueness
forall shape u_unused u.
ArrayShape shape =>
TypeBase shape u_unused -> shape -> u -> TypeBase shape u
arrayOf (Param (TypeBase Shape Uniqueness) -> TypeBase Shape Uniqueness
forall dec. DeclTyped dec => Param dec -> TypeBase Shape Uniqueness
paramDeclType Param (TypeBase Shape Uniqueness)
merge_param) ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w]) Uniqueness
Unique
        SubExp
expanded_init <- [Char] -> SubExp -> m SubExp
expandedInit [Char]
param_name SubExp
merge_init
        (Param (TypeBase Shape Uniqueness), SubExp)
-> m (Param (TypeBase Shape Uniqueness), SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param (TypeBase Shape Uniqueness)
expanded_param, SubExp
expanded_init)
        where
          param_name :: [Char]
param_name = VName -> [Char]
baseString (VName -> [Char]) -> VName -> [Char]
forall a b. (a -> b) -> a -> b
$ Param (TypeBase Shape Uniqueness) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase Shape Uniqueness)
merge_param

      expandPatElem :: PatElem Type -> PatElem Type
expandPatElem (PatElem VName
name Type
t) =
        VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
name (Type -> PatElem Type) -> Type -> PatElem Type
forall a b. (a -> b) -> a -> b
$ Type -> SubExp -> Type
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
arrayOfRow Type
t SubExp
w

-- We need to copy some initial arguments because otherwise the result
-- of the loop might alias the input (if the number of iterations is
-- 0), which is a problem if the result is consumed.
maybeCopyInitial ::
  (MonadBuilder m) =>
  (VName -> Bool) ->
  SeqLoop ->
  m SeqLoop
maybeCopyInitial :: (VName -> Bool) -> SeqLoop -> m SeqLoop
maybeCopyInitial VName -> Bool
isMapInput (SeqLoop [Int]
perm Pat Type
loop_pat [(FParam SOACS, SubExp)]
merge LoopForm SOACS
form Body SOACS
body) =
  [Int]
-> Pat Type
-> [(FParam SOACS, SubExp)]
-> LoopForm SOACS
-> Body SOACS
-> SeqLoop
SeqLoop [Int]
perm Pat Type
loop_pat ([(Param (TypeBase Shape Uniqueness), SubExp)]
 -> LoopForm SOACS -> Body SOACS -> SeqLoop)
-> m [(Param (TypeBase Shape Uniqueness), SubExp)]
-> m (LoopForm SOACS -> Body SOACS -> SeqLoop)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Param (TypeBase Shape Uniqueness), SubExp)
 -> m (Param (TypeBase Shape Uniqueness), SubExp))
-> [(Param (TypeBase Shape Uniqueness), SubExp)]
-> m [(Param (TypeBase Shape Uniqueness), SubExp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Param (TypeBase Shape Uniqueness), SubExp)
-> m (Param (TypeBase Shape Uniqueness), SubExp)
f [(Param (TypeBase Shape Uniqueness), SubExp)]
[(FParam SOACS, SubExp)]
merge m (LoopForm SOACS -> Body SOACS -> SeqLoop)
-> m (LoopForm SOACS) -> m (Body SOACS -> SeqLoop)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> LoopForm SOACS -> m (LoopForm SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure LoopForm SOACS
form m (Body SOACS -> SeqLoop) -> m (Body SOACS) -> m SeqLoop
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Body SOACS -> m (Body SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Body SOACS
body
  where
    f :: (Param (TypeBase Shape Uniqueness), SubExp)
-> m (Param (TypeBase Shape Uniqueness), SubExp)
f (Param (TypeBase Shape Uniqueness)
p, Var VName
arg)
      | VName -> Bool
isMapInput VName
arg,
        Array {} <- Param (TypeBase Shape Uniqueness) -> Type
forall dec. Typed dec => Param dec -> Type
paramType Param (TypeBase Shape Uniqueness)
p =
          (Param (TypeBase Shape Uniqueness)
p,) (SubExp -> (Param (TypeBase Shape Uniqueness), SubExp))
-> m SubExp -> m (Param (TypeBase Shape Uniqueness), SubExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> Exp (Rep m) -> m SubExp
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m SubExp
letSubExp (VName -> [Char]
baseString (Param (TypeBase Shape Uniqueness) -> VName
forall dec. Param dec -> VName
paramName Param (TypeBase Shape Uniqueness)
p) [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"_inter_copy") (BasicOp -> Exp (Rep m)
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
forall a b. (a -> b) -> a -> b
$ VName -> BasicOp
Copy VName
arg)
    f (Param (TypeBase Shape Uniqueness)
p, SubExp
arg) =
      (Param (TypeBase Shape Uniqueness), SubExp)
-> m (Param (TypeBase Shape Uniqueness), SubExp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param (TypeBase Shape Uniqueness)
p, SubExp
arg)

manifestMaps :: [LoopNesting] -> [VName] -> Stms SOACS -> ([VName], Stms SOACS)
manifestMaps :: [LoopNesting] -> [VName] -> Stms SOACS -> ([VName], Stms SOACS)
manifestMaps [] [VName]
res Stms SOACS
stms = ([VName]
res, Stms SOACS
stms)
manifestMaps (LoopNesting
n : [LoopNesting]
ns) [VName]
res Stms SOACS
stms =
  let ([VName]
res', Stms SOACS
stms') = [LoopNesting] -> [VName] -> Stms SOACS -> ([VName], Stms SOACS)
manifestMaps [LoopNesting]
ns [VName]
res Stms SOACS
stms
      ([Param Type]
params, [VName]
arrs) = [(Param Type, VName)] -> ([Param Type], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Param Type, VName)] -> ([Param Type], [VName]))
-> [(Param Type, VName)] -> ([Param Type], [VName])
forall a b. (a -> b) -> a -> b
$ LoopNesting -> [(Param Type, VName)]
loopNestingParamsAndArrs LoopNesting
n
      lam :: Lambda SOACS
lam =
        [LParam SOACS] -> Body SOACS -> [Type] -> Lambda SOACS
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda
          [Param Type]
[LParam SOACS]
params
          (Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms SOACS
stms' (Result -> Body SOACS) -> Result -> Body SOACS
forall a b. (a -> b) -> a -> b
$ [VName] -> Result
varsRes [VName]
res')
          ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall u. TypeBase Shape u -> TypeBase Shape u
rowType ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes (LoopNesting -> Pat Type
loopNestingPat LoopNesting
n))
   in ( Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames (Pat Type -> [VName]) -> Pat Type -> [VName]
forall a b. (a -> b) -> a -> b
$ LoopNesting -> Pat Type
loopNestingPat LoopNesting
n,
        Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm (Stm SOACS -> Stms SOACS) -> Stm SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$
          Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let (LoopNesting -> Pat Type
loopNestingPat LoopNesting
n) (LoopNesting -> StmAux ()
loopNestingAux LoopNesting
n) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$
            Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma (LoopNesting -> SubExp
loopNestingWidth LoopNesting
n) [VName]
arrs (Lambda SOACS -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda SOACS
lam)
      )

-- | Given a (parallel) map nesting and an inner sequential loop, move
-- the maps inside the sequential loop.  The result is several
-- statements - one of these will be the loop, which will then contain
-- statements with @map@ expressions.
interchangeLoops ::
  (MonadFreshNames m, HasScope SOACS m) =>
  KernelNest ->
  SeqLoop ->
  m (Stms SOACS)
interchangeLoops :: KernelNest -> SeqLoop -> m (Stms SOACS)
interchangeLoops KernelNest
full_nest = [LoopNesting] -> SeqLoop -> m (Stms SOACS)
recurse (KernelNest -> [LoopNesting]
kernelNestLoops KernelNest
full_nest)
  where
    recurse :: [LoopNesting] -> SeqLoop -> m (Stms SOACS)
recurse [LoopNesting]
nest SeqLoop
loop
      | ([LoopNesting]
ns, [LoopNesting
n]) <- Int -> [LoopNesting] -> ([LoopNesting], [LoopNesting])
forall a. Int -> [a] -> ([a], [a])
splitFromEnd Int
1 [LoopNesting]
nest = do
          let isMapParameter :: VName -> Maybe VName
isMapParameter VName
v =
                (Param Type, VName) -> VName
forall a b. (a, b) -> b
snd ((Param Type, VName) -> VName)
-> Maybe (Param Type, VName) -> Maybe VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Param Type, VName) -> Bool)
-> [(Param Type, VName)] -> Maybe (Param Type, VName)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v) (VName -> Bool)
-> ((Param Type, VName) -> VName) -> (Param Type, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> VName)
-> ((Param Type, VName) -> Param Type)
-> (Param Type, VName)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param Type, VName) -> Param Type
forall a b. (a, b) -> a
fst) (LoopNesting -> [(Param Type, VName)]
loopNestingParamsAndArrs LoopNesting
n)
              isMapInput :: VName -> Bool
isMapInput VName
v =
                VName
v VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ((Param Type, VName) -> VName) -> [(Param Type, VName)] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Param Type, VName) -> VName
forall a b. (a, b) -> b
snd (LoopNesting -> [(Param Type, VName)]
loopNestingParamsAndArrs LoopNesting
n)
          (SeqLoop
loop', Stms SOACS
stms) <-
            Builder SOACS SeqLoop -> m (SeqLoop, Stms SOACS)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder SOACS SeqLoop -> m (SeqLoop, Stms SOACS))
-> (Builder SOACS SeqLoop -> Builder SOACS SeqLoop)
-> Builder SOACS SeqLoop
-> m (SeqLoop, Stms SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope SOACS -> Builder SOACS SeqLoop -> Builder SOACS SeqLoop
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (KernelNest -> Scope SOACS
forall rep. (LParamInfo rep ~ Type) => KernelNest -> Scope rep
scopeOfKernelNest KernelNest
full_nest) (Builder SOACS SeqLoop -> m (SeqLoop, Stms SOACS))
-> Builder SOACS SeqLoop -> m (SeqLoop, Stms SOACS)
forall a b. (a -> b) -> a -> b
$
              (VName -> Bool) -> SeqLoop -> Builder SOACS SeqLoop
forall (m :: * -> *).
MonadBuilder m =>
(VName -> Bool) -> SeqLoop -> m SeqLoop
maybeCopyInitial VName -> Bool
isMapInput
                (SeqLoop -> Builder SOACS SeqLoop)
-> Builder SOACS SeqLoop -> Builder SOACS SeqLoop
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (VName -> Maybe VName)
-> SeqLoop -> LoopNesting -> Builder SOACS SeqLoop
forall (m :: * -> *).
(MonadBuilder m, LocalScope SOACS m) =>
(VName -> Maybe VName) -> SeqLoop -> LoopNesting -> m SeqLoop
interchangeLoop VName -> Maybe VName
isMapParameter SeqLoop
loop LoopNesting
n

          -- Only safe to continue interchanging if we didn't need to add
          -- any new statements; otherwise we manifest the remaining nests
          -- as Maps and hand them back to the flattener.
          if Stms SOACS -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Stms SOACS
stms
            then [LoopNesting] -> SeqLoop -> m (Stms SOACS)
recurse [LoopNesting]
ns SeqLoop
loop'
            else
              let loop_stm :: Stm SOACS
loop_stm = SeqLoop -> Stm SOACS
seqLoopStm SeqLoop
loop'
                  names :: [VName]
names = [Int] -> [VName] -> [VName]
forall a. [Int] -> [a] -> [a]
rearrangeShape (SeqLoop -> [Int]
loopPerm SeqLoop
loop') (Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames (Stm SOACS -> Pat (LetDec SOACS)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat Stm SOACS
loop_stm))
               in Stms SOACS -> m (Stms SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS -> m (Stms SOACS)) -> Stms SOACS -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ ([VName], Stms SOACS) -> Stms SOACS
forall a b. (a, b) -> b
snd (([VName], Stms SOACS) -> Stms SOACS)
-> ([VName], Stms SOACS) -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ [LoopNesting] -> [VName] -> Stms SOACS -> ([VName], Stms SOACS)
manifestMaps [LoopNesting]
ns [VName]
names (Stms SOACS -> ([VName], Stms SOACS))
-> Stms SOACS -> ([VName], Stms SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS
stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
loop_stm
      | Bool
otherwise = Stms SOACS -> m (Stms SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS -> m (Stms SOACS)) -> Stms SOACS -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm (Stm SOACS -> Stms SOACS) -> Stm SOACS -> Stms SOACS
forall a b. (a -> b) -> a -> b
$ SeqLoop -> Stm SOACS
seqLoopStm SeqLoop
loop

-- | An encoding of a branch with alongside its result pattern.
data Branch
  = Branch [Int] (Pat Type) SubExp (Body SOACS) (Body SOACS) (IfDec (BranchType SOACS))

branchStm :: Branch -> Stm SOACS
branchStm :: Branch -> Stm SOACS
branchStm (Branch [Int]
_ Pat Type
pat SubExp
cond Body SOACS
tbranch Body SOACS
fbranch IfDec (BranchType SOACS)
ret) =
  Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat Type
Pat (LetDec SOACS)
pat (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$ SubExp
-> Body SOACS
-> Body SOACS
-> IfDec (BranchType SOACS)
-> Exp SOACS
forall rep.
SubExp -> Body rep -> Body rep -> IfDec (BranchType rep) -> Exp rep
If SubExp
cond Body SOACS
tbranch Body SOACS
fbranch IfDec (BranchType SOACS)
ret

interchangeBranch1 ::
  (MonadBuilder m) =>
  Branch ->
  LoopNesting ->
  m Branch
interchangeBranch1 :: Branch -> LoopNesting -> m Branch
interchangeBranch1
  (Branch [Int]
perm Pat Type
branch_pat SubExp
cond Body SOACS
tbranch Body SOACS
fbranch (IfDec [BranchType SOACS]
ret IfSort
if_sort))
  (MapNesting Pat Type
pat StmAux ()
aux SubExp
w [(Param Type, VName)]
params_and_arrs) = do
    let ret' :: [TypeBase (ShapeBase (Ext SubExp)) NoUniqueness]
ret' = (TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
 -> TypeBase (ShapeBase (Ext SubExp)) NoUniqueness)
-> [TypeBase (ShapeBase (Ext SubExp)) NoUniqueness]
-> [TypeBase (ShapeBase (Ext SubExp)) NoUniqueness]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
-> Ext SubExp -> TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` SubExp -> Ext SubExp
forall a. a -> Ext a
Free SubExp
w) [TypeBase (ShapeBase (Ext SubExp)) NoUniqueness]
[BranchType SOACS]
ret
        pat' :: Pat Type
pat' = [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ [Int] -> [PatElem Type] -> [PatElem Type]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm ([PatElem Type] -> [PatElem Type])
-> [PatElem Type] -> [PatElem Type]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
pat

        ([Param Type]
params, [VName]
arrs) = [(Param Type, VName)] -> ([Param Type], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param Type, VName)]
params_and_arrs
        lam_ret :: [Type]
lam_ret = [Int] -> [Type] -> [Type]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall u. TypeBase Shape u -> TypeBase Shape u
rowType ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [Type]
forall dec. Typed dec => Pat dec -> [Type]
patTypes Pat Type
pat

        branch_pat' :: Pat Type
branch_pat' =
          [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ (PatElem Type -> PatElem Type) -> [PatElem Type] -> [PatElem Type]
forall a b. (a -> b) -> [a] -> [b]
map ((Type -> Type) -> PatElem Type -> PatElem Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type -> SubExp -> Type
forall d.
ArrayShape (ShapeBase d) =>
TypeBase (ShapeBase d) NoUniqueness
-> d -> TypeBase (ShapeBase d) NoUniqueness
`arrayOfRow` SubExp
w)) ([PatElem Type] -> [PatElem Type])
-> [PatElem Type] -> [PatElem Type]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
branch_pat

        mkBranch :: Body SOACS -> m (Body SOACS)
mkBranch Body SOACS
branch = (Body SOACS -> m (Body SOACS)
forall rep (m :: * -> *).
(Renameable rep, MonadFreshNames m) =>
Body rep -> m (Body rep)
renameBody (Body SOACS -> m (Body SOACS)) -> m (Body SOACS) -> m (Body SOACS)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) (m (Body SOACS) -> m (Body SOACS))
-> m (Body SOACS) -> m (Body SOACS)
forall a b. (a -> b) -> a -> b
$ do
          let lam :: Lambda SOACS
lam = [LParam SOACS] -> Body SOACS -> [Type] -> Lambda SOACS
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda [Param Type]
[LParam SOACS]
params Body SOACS
branch [Type]
lam_ret
              res :: Result
res = [VName] -> Result
varsRes ([VName] -> Result) -> [VName] -> Result
forall a b. (a -> b) -> a -> b
$ Pat Type -> [VName]
forall dec. Pat dec -> [VName]
patNames Pat Type
branch_pat'
              map_stm :: Stm SOACS
map_stm = Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat Type
Pat (LetDec SOACS)
branch_pat' StmAux ()
StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$ Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w [VName]
arrs (ScremaForm SOACS -> SOAC SOACS) -> ScremaForm SOACS -> SOAC SOACS
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda SOACS
lam
          Body SOACS -> m (Body SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body SOACS -> m (Body SOACS)) -> Body SOACS -> m (Body SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS -> Result -> Body SOACS
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody (Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm Stm SOACS
map_stm) Result
res

    Body SOACS
tbranch' <- Body SOACS -> m (Body SOACS)
mkBranch Body SOACS
tbranch
    Body SOACS
fbranch' <- Body SOACS -> m (Body SOACS)
mkBranch Body SOACS
fbranch
    Branch -> m Branch
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Branch -> m Branch) -> Branch -> m Branch
forall a b. (a -> b) -> a -> b
$
      [Int]
-> Pat Type
-> SubExp
-> Body SOACS
-> Body SOACS
-> IfDec (BranchType SOACS)
-> Branch
Branch [Int
0 .. Pat Type -> Int
forall dec. Pat dec -> Int
patSize Pat Type
pat Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] Pat Type
pat' SubExp
cond Body SOACS
tbranch' Body SOACS
fbranch' (IfDec (BranchType SOACS) -> Branch)
-> IfDec (BranchType SOACS) -> Branch
forall a b. (a -> b) -> a -> b
$
        [TypeBase (ShapeBase (Ext SubExp)) NoUniqueness]
-> IfSort -> IfDec (TypeBase (ShapeBase (Ext SubExp)) NoUniqueness)
forall rt. [rt] -> IfSort -> IfDec rt
IfDec [TypeBase (ShapeBase (Ext SubExp)) NoUniqueness]
ret' IfSort
if_sort

interchangeBranch ::
  (MonadFreshNames m, HasScope SOACS m) =>
  KernelNest ->
  Branch ->
  m (Stms SOACS)
interchangeBranch :: KernelNest -> Branch -> m (Stms SOACS)
interchangeBranch KernelNest
nest Branch
loop = do
  (Branch
loop', Stms SOACS
stms) <-
    Builder SOACS Branch -> m (Branch, Stms SOACS)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder SOACS Branch -> m (Branch, Stms SOACS))
-> Builder SOACS Branch -> m (Branch, Stms SOACS)
forall a b. (a -> b) -> a -> b
$ (Branch -> LoopNesting -> Builder SOACS Branch)
-> Branch -> [LoopNesting] -> Builder SOACS Branch
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Branch -> LoopNesting -> Builder SOACS Branch
forall (m :: * -> *).
MonadBuilder m =>
Branch -> LoopNesting -> m Branch
interchangeBranch1 Branch
loop ([LoopNesting] -> Builder SOACS Branch)
-> [LoopNesting] -> Builder SOACS Branch
forall a b. (a -> b) -> a -> b
$ [LoopNesting] -> [LoopNesting]
forall a. [a] -> [a]
reverse ([LoopNesting] -> [LoopNesting]) -> [LoopNesting] -> [LoopNesting]
forall a b. (a -> b) -> a -> b
$ KernelNest -> [LoopNesting]
kernelNestLoops KernelNest
nest
  Stms SOACS -> m (Stms SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS -> m (Stms SOACS)) -> Stms SOACS -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS
stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm (Branch -> Stm SOACS
branchStm Branch
loop')

-- | An encoding of a WithAcc with alongside its result pattern.
data WithAccStm
  = WithAccStm [Int] (Pat Type) [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] (Lambda SOACS)

withAccStm :: WithAccStm -> Stm SOACS
withAccStm :: WithAccStm -> Stm SOACS
withAccStm (WithAccStm [Int]
_ Pat Type
pat [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs Lambda SOACS
lam) =
  Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat Type
Pat (LetDec SOACS)
pat (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
forall a b. (a -> b) -> a -> b
$ [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
-> Lambda SOACS -> Exp SOACS
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs Lambda SOACS
lam

interchangeWithAcc1 ::
  (MonadBuilder m, Rep m ~ SOACS) =>
  WithAccStm ->
  LoopNesting ->
  m WithAccStm
interchangeWithAcc1 :: WithAccStm -> LoopNesting -> m WithAccStm
interchangeWithAcc1
  (WithAccStm [Int]
perm Pat Type
_withacc_pat [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs Lambda SOACS
acc_lam)
  (MapNesting Pat Type
map_pat StmAux ()
map_aux SubExp
w [(Param Type, VName)]
params_and_arrs) = do
    [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs' <- ((Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
 -> m (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
-> [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
-> m [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
-> m (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
onInput [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs
    [Param Type]
lam_params' <- [Param Type] -> m [Param Type]
forall (m :: * -> *) a.
MonadFreshNames m =>
[Param a] -> m [Param a]
newAccLamParams ([Param Type] -> m [Param Type]) -> [Param Type] -> m [Param Type]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda SOACS
acc_lam
    Param Type
iota_p <- [Char] -> Type -> m (Param Type)
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"iota_p" (Type -> m (Param Type)) -> Type -> m (Param Type)
forall a b. (a -> b) -> a -> b
$ PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64
    Lambda SOACS
acc_lam' <- SubExp -> Lambda SOACS -> m (Lambda SOACS)
trLam (VName -> SubExp
Var (Param Type -> VName
forall dec. Param dec -> VName
paramName Param Type
iota_p)) (Lambda SOACS -> m (Lambda SOACS))
-> (m Result -> m (Lambda SOACS)) -> m Result -> m (Lambda SOACS)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< [LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
forall (m :: * -> *).
MonadBuilder m =>
[LParam (Rep m)] -> m Result -> m (Lambda (Rep m))
mkLambda [Param Type]
[LParam (Rep m)]
lam_params' (m Result -> m (Lambda SOACS)) -> m Result -> m (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$ do
      let acc_params :: [Param Type]
acc_params = Int -> [Param Type] -> [Param Type]
forall a. Int -> [a] -> [a]
drop ([(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs) [Param Type]
lam_params'
          orig_acc_params :: [Param Type]
orig_acc_params = Int -> [Param Type] -> [Param Type]
forall a. Int -> [a] -> [a]
drop ([(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs) ([Param Type] -> [Param Type]) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda SOACS
acc_lam
      VName
iota_w <-
        [Char] -> Exp (Rep m) -> m VName
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m VName
letExp [Char]
"acc_inter_iota" (Exp SOACS -> m VName)
-> (BasicOp -> Exp SOACS) -> BasicOp -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> m VName) -> BasicOp -> m VName
forall a b. (a -> b) -> a -> b
$
          SubExp -> SubExp -> SubExp -> IntType -> BasicOp
Iota SubExp
w (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
0) (IntType -> Integer -> SubExp
intConst IntType
Int64 Integer
1) IntType
Int64
      let ([Param Type]
params, [VName]
arrs) = [(Param Type, VName)] -> ([Param Type], [VName])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Param Type, VName)]
params_and_arrs
          maplam_ret :: [Type]
maplam_ret = Lambda SOACS -> [Type]
forall rep. Lambda rep -> [Type]
lambdaReturnType Lambda SOACS
acc_lam
          maplam :: Lambda SOACS
maplam = [LParam SOACS] -> Body SOACS -> [Type] -> Lambda SOACS
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda (Param Type
iota_p Param Type -> [Param Type] -> [Param Type]
forall a. a -> [a] -> [a]
: [Param Type]
orig_acc_params [Param Type] -> [Param Type] -> [Param Type]
forall a. [a] -> [a] -> [a]
++ [Param Type]
params) (Lambda SOACS -> Body SOACS
forall rep. Lambda rep -> Body rep
lambdaBody Lambda SOACS
acc_lam) [Type]
maplam_ret
      StmAux () -> m Result -> m Result
forall (m :: * -> *) anyrep a.
MonadBuilder m =>
StmAux anyrep -> m a -> m a
auxing StmAux ()
map_aux (m Result -> m Result)
-> (Exp SOACS -> m Result) -> Exp SOACS -> m Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([SubExp] -> Result) -> m [SubExp] -> m Result
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [SubExp] -> Result
subExpsRes (m [SubExp] -> m Result)
-> (Exp SOACS -> m [SubExp]) -> Exp SOACS -> m Result
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Exp (Rep m) -> m [SubExp]
forall (m :: * -> *).
MonadBuilder m =>
[Char] -> Exp (Rep m) -> m [SubExp]
letTupExp' [Char]
"withacc_inter" (Exp SOACS -> m Result) -> Exp SOACS -> m Result
forall a b. (a -> b) -> a -> b
$
        Op SOACS -> Exp SOACS
forall rep. Op rep -> Exp rep
Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
Screma SubExp
w (VName
iota_w VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: (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]
acc_params [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
arrs) (Lambda SOACS -> ScremaForm SOACS
forall rep. Lambda rep -> ScremaForm rep
mapSOAC Lambda SOACS
maplam)
    let pat :: Pat Type
pat = [PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat ([PatElem Type] -> Pat Type) -> [PatElem Type] -> Pat Type
forall a b. (a -> b) -> a -> b
$ [Int] -> [PatElem Type] -> [PatElem Type]
forall a. [Int] -> [a] -> [a]
rearrangeShape [Int]
perm ([PatElem Type] -> [PatElem Type])
-> [PatElem Type] -> [PatElem Type]
forall a b. (a -> b) -> a -> b
$ Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems Pat Type
map_pat
    WithAccStm -> m WithAccStm
forall (f :: * -> *) a. Applicative f => a -> f a
pure (WithAccStm -> m WithAccStm) -> WithAccStm -> m WithAccStm
forall a b. (a -> b) -> a -> b
$ [Int]
-> Pat Type
-> [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
-> Lambda SOACS
-> WithAccStm
WithAccStm [Int]
perm Pat Type
pat [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs' Lambda SOACS
acc_lam'
    where
      newAccLamParams :: [Param a] -> m [Param a]
newAccLamParams [Param a]
ps = do
        let ([Param a]
cert_ps, [Param a]
acc_ps) = Int -> [Param a] -> ([Param a], [Param a])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Param a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Param a]
ps Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) [Param a]
ps
        -- Should not rename the certificates.
        [Param a]
acc_ps' <- [Param a] -> (Param a -> m (Param a)) -> m [Param a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Param a]
acc_ps ((Param a -> m (Param a)) -> m [Param a])
-> (Param a -> m (Param a)) -> m [Param a]
forall a b. (a -> b) -> a -> b
$ \(Param Attrs
attrs VName
v a
t) ->
          Attrs -> VName -> a -> Param a
forall dec. Attrs -> VName -> dec -> Param dec
Param Attrs
attrs (VName -> a -> Param a) -> m VName -> m (a -> Param a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName (VName -> [Char]
baseString VName
v) m (a -> Param a) -> m a -> m (Param a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
t
        [Param a] -> m [Param a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Param a] -> m [Param a]) -> [Param a] -> m [Param a]
forall a b. (a -> b) -> a -> b
$ [Param a]
cert_ps [Param a] -> [Param a] -> [Param a]
forall a. Semigroup a => a -> a -> a
<> [Param a]
acc_ps'

      num_accs :: Int
num_accs = [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
inputs
      acc_certs :: [VName]
acc_certs = (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] -> [VName]) -> [Param Type] -> [VName]
forall a b. (a -> b) -> a -> b
$ Int -> [Param Type] -> [Param Type]
forall a. Int -> [a] -> [a]
take Int
num_accs ([Param Type] -> [Param Type]) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> [LParam SOACS]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda SOACS
acc_lam
      onArr :: VName -> m VName
onArr VName
v =
        VName -> m VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> m VName)
-> (Maybe (Param Type, VName) -> VName)
-> Maybe (Param Type, VName)
-> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName
-> ((Param Type, VName) -> VName)
-> Maybe (Param Type, VName)
-> VName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe VName
v (Param Type, VName) -> VName
forall a b. (a, b) -> b
snd (Maybe (Param Type, VName) -> m VName)
-> Maybe (Param Type, VName) -> m VName
forall a b. (a -> b) -> a -> b
$
          ((Param Type, VName) -> Bool)
-> [(Param Type, VName)] -> Maybe (Param Type, VName)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v) (VName -> Bool)
-> ((Param Type, VName) -> VName) -> (Param Type, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> VName
forall dec. Param dec -> VName
paramName (Param Type -> VName)
-> ((Param Type, VName) -> Param Type)
-> (Param Type, VName)
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param Type, VName) -> Param Type
forall a b. (a, b) -> a
fst) [(Param Type, VName)]
params_and_arrs
      onInput :: (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
-> m (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
onInput (Shape
shape, [VName]
arrs, Maybe (Lambda SOACS, [SubExp])
op) =
        ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
shape,,) ([VName]
 -> Maybe (Lambda SOACS, [SubExp])
 -> (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
-> m [VName]
-> m (Maybe (Lambda SOACS, [SubExp])
      -> (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> m VName) -> [VName] -> m [VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM VName -> m VName
onArr [VName]
arrs m (Maybe (Lambda SOACS, [SubExp])
   -> (Shape, [VName], Maybe (Lambda SOACS, [SubExp])))
-> m (Maybe (Lambda SOACS, [SubExp]))
-> m (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((Lambda SOACS, [SubExp]) -> m (Lambda SOACS, [SubExp]))
-> Maybe (Lambda SOACS, [SubExp])
-> m (Maybe (Lambda SOACS, [SubExp]))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Lambda SOACS, [SubExp]) -> m (Lambda SOACS, [SubExp])
forall (m :: * -> *) rep shape u b.
(MonadFreshNames m, LParamInfo rep ~ TypeBase shape u) =>
(Lambda rep, b) -> m (Lambda rep, b)
onOp Maybe (Lambda SOACS, [SubExp])
op

      onOp :: (Lambda rep, b) -> m (Lambda rep, b)
onOp (Lambda rep
op_lam, b
nes) = do
        -- We need to add an additional index parameter because we are
        -- extending the index space of the accumulator.
        Param (TypeBase shape u)
idx_p <- [Char] -> TypeBase shape u -> m (Param (TypeBase shape u))
forall (m :: * -> *) dec.
MonadFreshNames m =>
[Char] -> dec -> m (Param dec)
newParam [Char]
"idx" (TypeBase shape u -> m (Param (TypeBase shape u)))
-> TypeBase shape u -> m (Param (TypeBase shape u))
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase shape u
forall shape u. PrimType -> TypeBase shape u
Prim PrimType
int64
        (Lambda rep, b) -> m (Lambda rep, b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Lambda rep
op_lam {lambdaParams :: [LParam rep]
lambdaParams = Param (TypeBase shape u)
idx_p Param (TypeBase shape u)
-> [Param (TypeBase shape u)] -> [Param (TypeBase shape u)]
forall a. a -> [a] -> [a]
: Lambda rep -> [LParam rep]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
op_lam}, b
nes)

      trType :: TypeBase shape u -> TypeBase shape u
      trType :: TypeBase shape u -> TypeBase shape u
trType (Acc VName
acc Shape
ispace [Type]
ts u
u)
        | VName
acc VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
acc_certs =
            VName -> Shape -> [Type] -> u -> TypeBase shape u
forall shape u. VName -> Shape -> [Type] -> u -> TypeBase shape u
Acc VName
acc ([SubExp] -> Shape
forall d. [d] -> ShapeBase d
Shape [SubExp
w] Shape -> Shape -> Shape
forall a. Semigroup a => a -> a -> a
<> Shape
ispace) [Type]
ts u
u
      trType TypeBase shape u
t = TypeBase shape u
t

      trParam :: Param (TypeBase shape u) -> Param (TypeBase shape u)
      trParam :: Param (TypeBase shape u) -> Param (TypeBase shape u)
trParam = (TypeBase shape u -> TypeBase shape u)
-> Param (TypeBase shape u) -> Param (TypeBase shape u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase shape u -> TypeBase shape u
forall shape u. TypeBase shape u -> TypeBase shape u
trType

      trLam :: SubExp -> Lambda SOACS -> m (Lambda SOACS)
trLam SubExp
i (Lambda [LParam SOACS]
params Body SOACS
body [Type]
ret) =
        Scope SOACS -> m (Lambda SOACS) -> m (Lambda SOACS)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param Type] -> Scope SOACS
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [Param Type]
[LParam SOACS]
params) (m (Lambda SOACS) -> m (Lambda SOACS))
-> m (Lambda SOACS) -> m (Lambda SOACS)
forall a b. (a -> b) -> a -> b
$
          [LParam SOACS] -> Body SOACS -> [Type] -> Lambda SOACS
forall rep. [LParam rep] -> Body rep -> [Type] -> Lambda rep
Lambda ((Param Type -> Param Type) -> [Param Type] -> [Param Type]
forall a b. (a -> b) -> [a] -> [b]
map Param Type -> Param Type
forall shape u.
Param (TypeBase shape u) -> Param (TypeBase shape u)
trParam [Param Type]
[LParam SOACS]
params) (Body SOACS -> [Type] -> Lambda SOACS)
-> m (Body SOACS) -> m ([Type] -> Lambda SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> Body SOACS -> m (Body SOACS)
trBody SubExp
i Body SOACS
body m ([Type] -> Lambda SOACS) -> m [Type] -> m (Lambda SOACS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Type] -> m [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
forall shape u. TypeBase shape u -> TypeBase shape u
trType [Type]
ret)

      trBody :: SubExp -> Body SOACS -> m (Body SOACS)
trBody SubExp
i (Body BodyDec SOACS
dec Stms SOACS
stms Result
res) =
        Stms SOACS -> m (Body SOACS) -> m (Body SOACS)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf Stms SOACS
stms (m (Body SOACS) -> m (Body SOACS))
-> m (Body SOACS) -> m (Body SOACS)
forall a b. (a -> b) -> a -> b
$ BodyDec SOACS -> Stms SOACS -> Result -> Body SOACS
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body BodyDec SOACS
dec (Stms SOACS -> Result -> Body SOACS)
-> m (Stms SOACS) -> m (Result -> Body SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Stm SOACS -> m (Stm SOACS)) -> Stms SOACS -> m (Stms SOACS)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (SubExp -> Stm SOACS -> m (Stm SOACS)
trStm SubExp
i) Stms SOACS
stms m (Result -> Body SOACS) -> m Result -> m (Body SOACS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Result -> m Result
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res

      trStm :: SubExp -> Stm SOACS -> m (Stm SOACS)
trStm SubExp
i (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux Exp SOACS
e) =
        Pat (LetDec SOACS)
-> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ((Type -> Type) -> Pat Type -> Pat Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Type
forall shape u. TypeBase shape u -> TypeBase shape u
trType Pat Type
Pat (LetDec SOACS)
pat) StmAux (ExpDec SOACS)
aux (Exp SOACS -> Stm SOACS) -> m (Exp SOACS) -> m (Stm SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> Exp SOACS -> m (Exp SOACS)
trExp SubExp
i Exp SOACS
e

      trSOAC :: SubExp -> SOAC SOACS -> m (SOAC SOACS)
trSOAC SubExp
i = SOACMapper SOACS SOACS m -> SOAC SOACS -> m (SOAC SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper SOACS SOACS m
mapper
        where
          mapper :: SOACMapper SOACS SOACS m
mapper =
            SOACMapper Any Any m
forall (m :: * -> *) rep. Monad m => SOACMapper rep rep m
identitySOACMapper {mapOnSOACLambda :: Lambda SOACS -> m (Lambda SOACS)
mapOnSOACLambda = SubExp -> Lambda SOACS -> m (Lambda SOACS)
trLam SubExp
i}

      trExp :: SubExp -> Exp SOACS -> m (Exp SOACS)
trExp SubExp
i (WithAcc [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
acc_inputs Lambda SOACS
lam) =
        [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
-> Lambda SOACS -> Exp SOACS
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))]
acc_inputs (Lambda SOACS -> Exp SOACS) -> m (Lambda SOACS) -> m (Exp SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SubExp -> Lambda SOACS -> m (Lambda SOACS)
trLam SubExp
i Lambda SOACS
lam
      trExp SubExp
i (BasicOp (UpdateAcc VName
acc [SubExp]
is [SubExp]
ses)) = do
        Type
acc_t <- VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
acc
        Exp SOACS -> m (Exp SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp SOACS -> m (Exp SOACS)) -> Exp SOACS -> m (Exp SOACS)
forall a b. (a -> b) -> a -> b
$ case Type
acc_t of
          Acc VName
cert Shape
_ [Type]
_ NoUniqueness
_
            | VName
cert VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
acc_certs ->
                BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc VName
acc (SubExp
i SubExp -> [SubExp] -> [SubExp]
forall a. a -> [a] -> [a]
: [SubExp]
is) [SubExp]
ses
          Type
_ ->
            BasicOp -> Exp SOACS
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Exp SOACS) -> BasicOp -> Exp SOACS
forall a b. (a -> b) -> a -> b
$ VName -> [SubExp] -> [SubExp] -> BasicOp
UpdateAcc VName
acc [SubExp]
is [SubExp]
ses
      trExp SubExp
i Exp SOACS
e = Mapper SOACS SOACS m -> Exp SOACS -> m (Exp SOACS)
forall (m :: * -> *) frep trep.
(Applicative m, Monad m) =>
Mapper frep trep m -> Exp frep -> m (Exp trep)
mapExpM Mapper SOACS SOACS m
mapper Exp SOACS
e
        where
          mapper :: Mapper SOACS SOACS m
mapper =
            Mapper Any Any m
forall (m :: * -> *) rep. Monad m => Mapper rep rep m
identityMapper
              { mapOnBody :: Scope SOACS -> Body SOACS -> m (Body SOACS)
mapOnBody = \Scope SOACS
scope -> Scope SOACS -> m (Body SOACS) -> m (Body SOACS)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope Scope SOACS
scope (m (Body SOACS) -> m (Body SOACS))
-> (Body SOACS -> m (Body SOACS)) -> Body SOACS -> m (Body SOACS)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> Body SOACS -> m (Body SOACS)
trBody SubExp
i,
                mapOnRetType :: RetType SOACS -> m (RetType SOACS)
mapOnRetType = TypeBase (ShapeBase (Ext SubExp)) Uniqueness
-> m (TypeBase (ShapeBase (Ext SubExp)) Uniqueness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (ShapeBase (Ext SubExp)) Uniqueness
 -> m (TypeBase (ShapeBase (Ext SubExp)) Uniqueness))
-> (TypeBase (ShapeBase (Ext SubExp)) Uniqueness
    -> TypeBase (ShapeBase (Ext SubExp)) Uniqueness)
-> TypeBase (ShapeBase (Ext SubExp)) Uniqueness
-> m (TypeBase (ShapeBase (Ext SubExp)) Uniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (ShapeBase (Ext SubExp)) Uniqueness
-> TypeBase (ShapeBase (Ext SubExp)) Uniqueness
forall shape u. TypeBase shape u -> TypeBase shape u
trType,
                mapOnBranchType :: BranchType SOACS -> m (BranchType SOACS)
mapOnBranchType = TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
-> m (TypeBase (ShapeBase (Ext SubExp)) NoUniqueness)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
 -> m (TypeBase (ShapeBase (Ext SubExp)) NoUniqueness))
-> (TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
    -> TypeBase (ShapeBase (Ext SubExp)) NoUniqueness)
-> TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
-> m (TypeBase (ShapeBase (Ext SubExp)) NoUniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
-> TypeBase (ShapeBase (Ext SubExp)) NoUniqueness
forall shape u. TypeBase shape u -> TypeBase shape u
trType,
                mapOnFParam :: FParam SOACS -> m (FParam SOACS)
mapOnFParam = Param (TypeBase Shape Uniqueness)
-> m (Param (TypeBase Shape Uniqueness))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param (TypeBase Shape Uniqueness)
 -> m (Param (TypeBase Shape Uniqueness)))
-> (Param (TypeBase Shape Uniqueness)
    -> Param (TypeBase Shape Uniqueness))
-> Param (TypeBase Shape Uniqueness)
-> m (Param (TypeBase Shape Uniqueness))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param (TypeBase Shape Uniqueness)
-> Param (TypeBase Shape Uniqueness)
forall shape u.
Param (TypeBase shape u) -> Param (TypeBase shape u)
trParam,
                mapOnLParam :: LParam SOACS -> m (LParam SOACS)
mapOnLParam = Param Type -> m (Param Type)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Param Type -> m (Param Type))
-> (Param Type -> Param Type) -> Param Type -> m (Param Type)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Param Type -> Param Type
forall shape u.
Param (TypeBase shape u) -> Param (TypeBase shape u)
trParam,
                mapOnOp :: Op SOACS -> m (Op SOACS)
mapOnOp = SubExp -> SOAC SOACS -> m (SOAC SOACS)
trSOAC SubExp
i
              }

interchangeWithAcc ::
  (MonadFreshNames m, HasScope SOACS m) =>
  KernelNest ->
  WithAccStm ->
  m (Stms SOACS)
interchangeWithAcc :: KernelNest -> WithAccStm -> m (Stms SOACS)
interchangeWithAcc KernelNest
nest WithAccStm
withacc = do
  (WithAccStm
withacc', Stms SOACS
stms) <-
    Builder SOACS WithAccStm -> m (WithAccStm, Stms SOACS)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder SOACS WithAccStm -> m (WithAccStm, Stms SOACS))
-> Builder SOACS WithAccStm -> m (WithAccStm, Stms SOACS)
forall a b. (a -> b) -> a -> b
$ (WithAccStm -> LoopNesting -> Builder SOACS WithAccStm)
-> WithAccStm -> [LoopNesting] -> Builder SOACS WithAccStm
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM WithAccStm -> LoopNesting -> Builder SOACS WithAccStm
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ SOACS) =>
WithAccStm -> LoopNesting -> m WithAccStm
interchangeWithAcc1 WithAccStm
withacc ([LoopNesting] -> Builder SOACS WithAccStm)
-> [LoopNesting] -> Builder SOACS WithAccStm
forall a b. (a -> b) -> a -> b
$ [LoopNesting] -> [LoopNesting]
forall a. [a] -> [a]
reverse ([LoopNesting] -> [LoopNesting]) -> [LoopNesting] -> [LoopNesting]
forall a b. (a -> b) -> a -> b
$ KernelNest -> [LoopNesting]
kernelNestLoops KernelNest
nest
  Stms SOACS -> m (Stms SOACS)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms SOACS -> m (Stms SOACS)) -> Stms SOACS -> m (Stms SOACS)
forall a b. (a -> b) -> a -> b
$ Stms SOACS
stms Stms SOACS -> Stms SOACS -> Stms SOACS
forall a. Semigroup a => a -> a -> a
<> Stm SOACS -> Stms SOACS
forall rep. Stm rep -> Stms rep
oneStm (WithAccStm -> Stm SOACS
withAccStm WithAccStm
withacc')