{-# LANGUAGE TypeFamilies #-}

-- | Extraction of parallelism from a SOACs program.  This generates
-- parallel constructs aimed at CPU execution, which in particular may
-- involve ad-hoc irregular nested parallelism.
module Futhark.Pass.ExtractMulticore (extractMulticore) where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Data.Bitraversable
import Futhark.IR
import Futhark.IR.MC
import Futhark.IR.MC qualified as MC
import Futhark.IR.SOACS hiding
  ( Body,
    Exp,
    LParam,
    Lambda,
    Pat,
    Stm,
  )
import Futhark.IR.SOACS qualified as SOACS
import Futhark.Pass
import Futhark.Pass.ExtractKernels.DistributeNests
import Futhark.Pass.ExtractKernels.ToGPU (injectSOACS)
import Futhark.Tools
import Futhark.Transform.Rename (Rename, renameSomething)
import Futhark.Util.Log

newtype ExtractM a = ExtractM (ReaderT (Scope MC) (State VNameSource) a)
  deriving
    ( (forall a b. (a -> b) -> ExtractM a -> ExtractM b)
-> (forall a b. a -> ExtractM b -> ExtractM a) -> Functor ExtractM
forall a b. a -> ExtractM b -> ExtractM a
forall a b. (a -> b) -> ExtractM a -> ExtractM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> ExtractM a -> ExtractM b
fmap :: forall a b. (a -> b) -> ExtractM a -> ExtractM b
$c<$ :: forall a b. a -> ExtractM b -> ExtractM a
<$ :: forall a b. a -> ExtractM b -> ExtractM a
Functor,
      Functor ExtractM
Functor ExtractM =>
(forall a. a -> ExtractM a)
-> (forall a b. ExtractM (a -> b) -> ExtractM a -> ExtractM b)
-> (forall a b c.
    (a -> b -> c) -> ExtractM a -> ExtractM b -> ExtractM c)
-> (forall a b. ExtractM a -> ExtractM b -> ExtractM b)
-> (forall a b. ExtractM a -> ExtractM b -> ExtractM a)
-> Applicative ExtractM
forall a. a -> ExtractM a
forall a b. ExtractM a -> ExtractM b -> ExtractM a
forall a b. ExtractM a -> ExtractM b -> ExtractM b
forall a b. ExtractM (a -> b) -> ExtractM a -> ExtractM b
forall a b c.
(a -> b -> c) -> ExtractM a -> ExtractM b -> ExtractM c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> ExtractM a
pure :: forall a. a -> ExtractM a
$c<*> :: forall a b. ExtractM (a -> b) -> ExtractM a -> ExtractM b
<*> :: forall a b. ExtractM (a -> b) -> ExtractM a -> ExtractM b
$cliftA2 :: forall a b c.
(a -> b -> c) -> ExtractM a -> ExtractM b -> ExtractM c
liftA2 :: forall a b c.
(a -> b -> c) -> ExtractM a -> ExtractM b -> ExtractM c
$c*> :: forall a b. ExtractM a -> ExtractM b -> ExtractM b
*> :: forall a b. ExtractM a -> ExtractM b -> ExtractM b
$c<* :: forall a b. ExtractM a -> ExtractM b -> ExtractM a
<* :: forall a b. ExtractM a -> ExtractM b -> ExtractM a
Applicative,
      Applicative ExtractM
Applicative ExtractM =>
(forall a b. ExtractM a -> (a -> ExtractM b) -> ExtractM b)
-> (forall a b. ExtractM a -> ExtractM b -> ExtractM b)
-> (forall a. a -> ExtractM a)
-> Monad ExtractM
forall a. a -> ExtractM a
forall a b. ExtractM a -> ExtractM b -> ExtractM b
forall a b. ExtractM a -> (a -> ExtractM b) -> ExtractM b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b. ExtractM a -> (a -> ExtractM b) -> ExtractM b
>>= :: forall a b. ExtractM a -> (a -> ExtractM b) -> ExtractM b
$c>> :: forall a b. ExtractM a -> ExtractM b -> ExtractM b
>> :: forall a b. ExtractM a -> ExtractM b -> ExtractM b
$creturn :: forall a. a -> ExtractM a
return :: forall a. a -> ExtractM a
Monad,
      HasScope MC,
      LocalScope MC,
      Monad ExtractM
ExtractM VNameSource
Monad ExtractM =>
ExtractM VNameSource
-> (VNameSource -> ExtractM ()) -> MonadFreshNames ExtractM
VNameSource -> ExtractM ()
forall (m :: * -> *).
Monad m =>
m VNameSource -> (VNameSource -> m ()) -> MonadFreshNames m
$cgetNameSource :: ExtractM VNameSource
getNameSource :: ExtractM VNameSource
$cputNameSource :: VNameSource -> ExtractM ()
putNameSource :: VNameSource -> ExtractM ()
MonadFreshNames
    )

-- XXX: throwing away the log here...
instance MonadLogger ExtractM where
  addLog :: Log -> ExtractM ()
addLog Log
_ = () -> ExtractM ()
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

indexArray :: VName -> LParam SOACS -> VName -> Stm MC
indexArray :: VName -> LParam SOACS -> VName -> Stm MC
indexArray VName
i (Param Attrs
_ VName
p LParamInfo SOACS
t) VName
arr =
  Pat (LetDec MC) -> StmAux (ExpDec MC) -> Exp MC -> Stm MC
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let ([PatElem Type] -> Pat Type
forall dec. [PatElem dec] -> Pat dec
Pat [VName -> Type -> PatElem Type
forall dec. VName -> dec -> PatElem dec
PatElem VName
p Type
LParamInfo SOACS
t]) (() -> StmAux ()
forall dec. dec -> StmAux dec
defAux ()) (Exp MC -> Stm MC) -> (BasicOp -> Exp MC) -> BasicOp -> Stm MC
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicOp -> Exp MC
forall rep. BasicOp -> Exp rep
BasicOp (BasicOp -> Stm MC) -> BasicOp -> Stm MC
forall a b. (a -> b) -> a -> b
$
    case LParamInfo SOACS
t of
      Acc {} -> SubExp -> BasicOp
SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ VName -> SubExp
Var VName
arr
      LParamInfo SOACS
_ -> VName -> Slice SubExp -> BasicOp
Index VName
arr (Slice SubExp -> BasicOp) -> Slice SubExp -> BasicOp
forall a b. (a -> b) -> a -> b
$ [DimIndex SubExp] -> Slice SubExp
forall d. [DimIndex d] -> Slice d
Slice ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (VName -> SubExp
Var VName
i) DimIndex SubExp -> [DimIndex SubExp] -> [DimIndex SubExp]
forall a. a -> [a] -> [a]
: (SubExp -> DimIndex SubExp) -> [SubExp] -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> DimIndex SubExp
sliceDim (Type -> [SubExp]
forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
arrayDims Type
LParamInfo SOACS
t)

mapLambdaToBody ::
  (Body SOACS -> ExtractM (Body MC)) ->
  VName ->
  Lambda SOACS ->
  [VName] ->
  ExtractM (Body MC)
mapLambdaToBody :: (Body SOACS -> ExtractM (Body MC))
-> VName -> Lambda SOACS -> [VName] -> ExtractM (Body MC)
mapLambdaToBody Body SOACS -> ExtractM (Body MC)
onBody VName
i Lambda SOACS
lam [VName]
arrs = do
  let indexings :: [Stm MC]
indexings = (Param Type -> VName -> Stm MC)
-> [Param Type] -> [VName] -> [Stm MC]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (VName -> LParam SOACS -> VName -> Stm MC
indexArray VName
i) (Lambda SOACS -> [LParam SOACS]
forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda SOACS
lam) [VName]
arrs
  Body () stms res <- [Stm MC] -> ExtractM (Body MC) -> ExtractM (Body MC)
forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf [Stm MC]
indexings (ExtractM (Body MC) -> ExtractM (Body MC))
-> ExtractM (Body MC) -> ExtractM (Body MC)
forall a b. (a -> b) -> a -> b
$ Body SOACS -> ExtractM (Body MC)
onBody (Body SOACS -> ExtractM (Body MC))
-> Body SOACS -> ExtractM (Body MC)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS -> Body SOACS
forall rep. Lambda rep -> Body rep
lambdaBody Lambda SOACS
lam
  pure $ Body () (stmsFromList indexings <> stms) res

mapLambdaToKernelBody ::
  (Body SOACS -> ExtractM (Body MC)) ->
  VName ->
  Lambda SOACS ->
  [VName] ->
  ExtractM (KernelBody MC)
mapLambdaToKernelBody :: (Body SOACS -> ExtractM (Body MC))
-> VName -> Lambda SOACS -> [VName] -> ExtractM (KernelBody MC)
mapLambdaToKernelBody Body SOACS -> ExtractM (Body MC)
onBody VName
i Lambda SOACS
lam [VName]
arrs = do
  Body () stms res <- (Body SOACS -> ExtractM (Body MC))
-> VName -> Lambda SOACS -> [VName] -> ExtractM (Body MC)
mapLambdaToBody Body SOACS -> ExtractM (Body MC)
onBody VName
i Lambda SOACS
lam [VName]
arrs
  let ret (SubExpRes Certs
cs SubExp
se) = ResultManifest -> Certs -> SubExp -> KernelResult
Returns ResultManifest
ResultMaySimplify Certs
cs SubExp
se
  pure $ KernelBody () stms $ map ret res

reduceToSegBinOp :: Reduce SOACS -> ExtractM (Stms MC, SegBinOp MC)
reduceToSegBinOp :: Reduce SOACS -> ExtractM (Stms MC, SegBinOp MC)
reduceToSegBinOp (Reduce Commutativity
comm Lambda SOACS
lam [SubExp]
nes) = do
  ((lam', nes', shape), stms) <- Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
-> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
 -> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC))
-> Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
-> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS
-> [SubExp]
-> Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
forall (m :: * -> *).
MonadBuilder m =>
Lambda SOACS
-> [SubExp] -> m (Lambda SOACS, [SubExp], ShapeBase SubExp)
determineReduceOp Lambda SOACS
lam [SubExp]
nes
  lam'' <- transformLambda lam'
  let comm'
        | Lambda SOACS -> Bool
forall rep. Lambda rep -> Bool
commutativeLambda Lambda SOACS
lam' = Commutativity
Commutative
        | Bool
otherwise = Commutativity
comm
  pure (stms, SegBinOp comm' lam'' nes' shape)

scanToSegBinOp :: Scan SOACS -> ExtractM (Stms MC, SegBinOp MC)
scanToSegBinOp :: Scan SOACS -> ExtractM (Stms MC, SegBinOp MC)
scanToSegBinOp (Scan Lambda SOACS
lam [SubExp]
nes) = do
  ((lam', nes', shape), stms) <- Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
-> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
 -> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC))
-> Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
-> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS
-> [SubExp]
-> Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
forall (m :: * -> *).
MonadBuilder m =>
Lambda SOACS
-> [SubExp] -> m (Lambda SOACS, [SubExp], ShapeBase SubExp)
determineReduceOp Lambda SOACS
lam [SubExp]
nes
  lam'' <- transformLambda lam'
  pure (stms, SegBinOp Noncommutative lam'' nes' shape)

histToSegBinOp :: SOACS.HistOp SOACS -> ExtractM (Stms MC, MC.HistOp MC)
histToSegBinOp :: HistOp SOACS -> ExtractM (Stms MC, HistOp MC)
histToSegBinOp (SOACS.HistOp ShapeBase SubExp
num_bins SubExp
rf [VName]
dests [SubExp]
nes Lambda SOACS
op) = do
  ((op', nes', shape), stms) <- Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
-> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
 -> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC))
-> Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
-> ExtractM ((Lambda SOACS, [SubExp], ShapeBase SubExp), Stms MC)
forall a b. (a -> b) -> a -> b
$ Lambda SOACS
-> [SubExp]
-> Builder MC (Lambda SOACS, [SubExp], ShapeBase SubExp)
forall (m :: * -> *).
MonadBuilder m =>
Lambda SOACS
-> [SubExp] -> m (Lambda SOACS, [SubExp], ShapeBase SubExp)
determineReduceOp Lambda SOACS
op [SubExp]
nes
  op'' <- transformLambda op'
  pure (stms, MC.HistOp num_bins rf dests nes' shape op'')

mkSegSpace :: (MonadFreshNames m) => SubExp -> m (VName, SegSpace)
mkSegSpace :: forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> m (VName, SegSpace)
mkSegSpace SubExp
w = do
  flat <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"flat_tid"
  gtid <- newVName "gtid"
  let space = VName -> [(VName, SubExp)] -> SegSpace
SegSpace VName
flat [(VName
gtid, SubExp
w)]
  pure (gtid, space)

transformStm :: Stm SOACS -> ExtractM (Stms MC)
transformStm :: Stm SOACS -> ExtractM (Stms MC)
transformStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (BasicOp BasicOp
op)) =
  Stms MC -> ExtractM (Stms MC)
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms MC -> ExtractM (Stms MC)) -> Stms MC -> ExtractM (Stms MC)
forall a b. (a -> b) -> a -> b
$ Stm MC -> Stms MC
forall rep. Stm rep -> Stms rep
oneStm (Stm MC -> Stms MC) -> Stm MC -> Stms MC
forall a b. (a -> b) -> a -> b
$ Pat (LetDec MC) -> StmAux (ExpDec MC) -> Exp MC -> Stm MC
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
Pat (LetDec MC)
pat StmAux (ExpDec SOACS)
StmAux (ExpDec MC)
aux (Exp MC -> Stm MC) -> Exp MC -> Stm MC
forall a b. (a -> b) -> a -> b
$ BasicOp -> Exp MC
forall rep. BasicOp -> Exp rep
BasicOp BasicOp
op
transformStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Apply Name
f [(SubExp, Diet)]
args [(RetType SOACS, RetAls)]
ret Safety
info)) =
  Stms MC -> ExtractM (Stms MC)
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stms MC -> ExtractM (Stms MC)) -> Stms MC -> ExtractM (Stms MC)
forall a b. (a -> b) -> a -> b
$ Stm MC -> Stms MC
forall rep. Stm rep -> Stms rep
oneStm (Stm MC -> Stms MC) -> Stm MC -> Stms MC
forall a b. (a -> b) -> a -> b
$ Pat (LetDec MC) -> StmAux (ExpDec MC) -> Exp MC -> Stm MC
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
Pat (LetDec MC)
pat StmAux (ExpDec SOACS)
StmAux (ExpDec MC)
aux (Exp MC -> Stm MC) -> Exp MC -> Stm MC
forall a b. (a -> b) -> a -> b
$ Name
-> [(SubExp, Diet)] -> [(RetType MC, RetAls)] -> Safety -> Exp MC
forall rep.
Name
-> [(SubExp, Diet)] -> [(RetType rep, RetAls)] -> Safety -> Exp rep
Apply Name
f [(SubExp, Diet)]
args [(RetType SOACS, RetAls)]
[(RetType MC, RetAls)]
ret Safety
info
transformStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Loop [(FParam SOACS, SubExp)]
merge LoopForm
form Body SOACS
body)) = do
  body' <-
    Scope MC -> ExtractM (Body MC) -> ExtractM (Body MC)
forall a. Scope MC -> ExtractM a -> ExtractM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (FParamInfo MC)] -> Scope MC
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (((Param (FParamInfo MC), SubExp) -> Param (FParamInfo MC))
-> [(Param (FParamInfo MC), SubExp)] -> [Param (FParamInfo MC)]
forall a b. (a -> b) -> [a] -> [b]
map (Param (FParamInfo MC), SubExp) -> Param (FParamInfo MC)
forall a b. (a, b) -> a
fst [(FParam SOACS, SubExp)]
[(Param (FParamInfo MC), SubExp)]
merge) Scope MC -> Scope MC -> Scope MC
forall a. Semigroup a => a -> a -> a
<> LoopForm -> Scope MC
forall rep. LoopForm -> Scope rep
scopeOfLoopForm LoopForm
form) (ExtractM (Body MC) -> ExtractM (Body MC))
-> ExtractM (Body MC) -> ExtractM (Body MC)
forall a b. (a -> b) -> a -> b
$
      Body SOACS -> ExtractM (Body MC)
transformBody Body SOACS
body
  pure $ oneStm $ Let pat aux $ Loop merge form body'
transformStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Match [SubExp]
ses [Case (Body SOACS)]
cases Body SOACS
defbody MatchDec (BranchType SOACS)
ret)) =
  Stm MC -> Stms MC
forall rep. Stm rep -> Stms rep
oneStm (Stm MC -> Stms MC) -> (Exp MC -> Stm MC) -> Exp MC -> Stms MC
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec MC) -> StmAux (ExpDec MC) -> Exp MC -> Stm MC
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
Pat (LetDec MC)
pat StmAux (ExpDec SOACS)
StmAux (ExpDec MC)
aux
    (Exp MC -> Stms MC) -> ExtractM (Exp MC) -> ExtractM (Stms MC)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([SubExp]
-> [Case (Body MC)]
-> Body MC
-> MatchDec (BranchType MC)
-> Exp MC
forall rep.
[SubExp]
-> [Case (Body rep)]
-> Body rep
-> MatchDec (BranchType rep)
-> Exp rep
Match [SubExp]
ses ([Case (Body MC)] -> Body MC -> MatchDec ExtType -> Exp MC)
-> ExtractM [Case (Body MC)]
-> ExtractM (Body MC -> MatchDec ExtType -> Exp MC)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Case (Body SOACS) -> ExtractM (Case (Body MC)))
-> [Case (Body SOACS)] -> ExtractM [Case (Body MC)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Case (Body SOACS) -> ExtractM (Case (Body MC))
transformCase [Case (Body SOACS)]
cases ExtractM (Body MC -> MatchDec ExtType -> Exp MC)
-> ExtractM (Body MC) -> ExtractM (MatchDec ExtType -> Exp MC)
forall a b. ExtractM (a -> b) -> ExtractM a -> ExtractM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Body SOACS -> ExtractM (Body MC)
transformBody Body SOACS
defbody ExtractM (MatchDec ExtType -> Exp MC)
-> ExtractM (MatchDec ExtType) -> ExtractM (Exp MC)
forall a b. ExtractM (a -> b) -> ExtractM a -> ExtractM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> MatchDec ExtType -> ExtractM (MatchDec ExtType)
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MatchDec ExtType
MatchDec (BranchType SOACS)
ret)
  where
    transformCase :: Case (Body SOACS) -> ExtractM (Case (Body MC))
transformCase (Case [Maybe PrimValue]
vs Body SOACS
body) = [Maybe PrimValue] -> Body MC -> Case (Body MC)
forall body. [Maybe PrimValue] -> body -> Case body
Case [Maybe PrimValue]
vs (Body MC -> Case (Body MC))
-> ExtractM (Body MC) -> ExtractM (Case (Body MC))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Body SOACS -> ExtractM (Body MC)
transformBody Body SOACS
body
transformStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (WithAcc [WithAccInput SOACS]
inputs Lambda SOACS
lam)) =
  Stm MC -> Stms MC
forall rep. Stm rep -> Stms rep
oneStm (Stm MC -> Stms MC) -> (Exp MC -> Stm MC) -> Exp MC -> Stms MC
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat (LetDec MC) -> StmAux (ExpDec MC) -> Exp MC -> Stm MC
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec SOACS)
Pat (LetDec MC)
pat StmAux (ExpDec SOACS)
StmAux (ExpDec MC)
aux
    (Exp MC -> Stms MC) -> ExtractM (Exp MC) -> ExtractM (Stms MC)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([WithAccInput MC] -> Lambda MC -> Exp MC
forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
WithAcc ([WithAccInput MC] -> Lambda MC -> Exp MC)
-> ExtractM [WithAccInput MC] -> ExtractM (Lambda MC -> Exp MC)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (WithAccInput SOACS -> ExtractM (WithAccInput MC))
-> [WithAccInput SOACS] -> ExtractM [WithAccInput MC]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM WithAccInput SOACS -> ExtractM (WithAccInput MC)
forall {t :: * -> *} {t :: * -> * -> *} {t} {t} {d}.
(Traversable t, Bitraversable t) =>
(t, t, t (t (Lambda SOACS) d))
-> ExtractM (t, t, t (t (Lambda MC) d))
transformInput [WithAccInput SOACS]
inputs ExtractM (Lambda MC -> Exp MC)
-> ExtractM (Lambda MC) -> ExtractM (Exp MC)
forall a b. ExtractM (a -> b) -> ExtractM a -> ExtractM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Lambda SOACS -> ExtractM (Lambda MC)
transformLambda Lambda SOACS
lam)
  where
    transformInput :: (t, t, t (t (Lambda SOACS) d))
-> ExtractM (t, t, t (t (Lambda MC) d))
transformInput (t
shape, t
arrs, t (t (Lambda SOACS) d)
op) =
      (t
shape,t
arrs,) (t (t (Lambda MC) d) -> (t, t, t (t (Lambda MC) d)))
-> ExtractM (t (t (Lambda MC) d))
-> ExtractM (t, t, t (t (Lambda MC) d))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (t (Lambda SOACS) d -> ExtractM (t (Lambda MC) d))
-> t (t (Lambda SOACS) d) -> ExtractM (t (t (Lambda MC) d))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b)
traverse ((Lambda SOACS -> ExtractM (Lambda MC))
-> (d -> ExtractM d)
-> t (Lambda SOACS) d
-> ExtractM (t (Lambda MC) d)
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse Lambda SOACS -> ExtractM (Lambda MC)
transformLambda d -> ExtractM d
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) t (t (Lambda SOACS) d)
op
transformStm (Let Pat (LetDec SOACS)
pat StmAux (ExpDec SOACS)
aux (Op Op SOACS
op)) =
  (Stm MC -> Stm MC) -> Stms MC -> Stms MC
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Certs -> Stm MC -> Stm MC
forall rep. Certs -> Stm rep -> Stm rep
certify (StmAux () -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts StmAux ()
StmAux (ExpDec SOACS)
aux)) (Stms MC -> Stms MC) -> ExtractM (Stms MC) -> ExtractM (Stms MC)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pat Type -> Attrs -> SOAC SOACS -> ExtractM (Stms MC)
transformSOAC Pat Type
Pat (LetDec SOACS)
pat (StmAux () -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs StmAux ()
StmAux (ExpDec SOACS)
aux) Op SOACS
SOAC SOACS
op

transformLambda :: Lambda SOACS -> ExtractM (Lambda MC)
transformLambda :: Lambda SOACS -> ExtractM (Lambda MC)
transformLambda (Lambda [LParam SOACS]
params [Type]
ret Body SOACS
body) =
  [LParam MC] -> [Type] -> Body MC -> Lambda MC
forall rep. [LParam rep] -> [Type] -> Body rep -> Lambda rep
Lambda [LParam SOACS]
[LParam MC]
params [Type]
ret
    (Body MC -> Lambda MC)
-> ExtractM (Body MC) -> ExtractM (Lambda MC)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Scope MC -> ExtractM (Body MC) -> ExtractM (Body MC)
forall a. Scope MC -> ExtractM a -> ExtractM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([LParam MC] -> Scope MC
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [LParam SOACS]
[LParam MC]
params) (Body SOACS -> ExtractM (Body MC)
transformBody Body SOACS
body)

transformStms :: Stms SOACS -> ExtractM (Stms MC)
transformStms :: Stms SOACS -> ExtractM (Stms MC)
transformStms Stms SOACS
stms =
  case Stms SOACS -> Maybe (Stm SOACS, Stms SOACS)
forall rep. Stms rep -> Maybe (Stm rep, Stms rep)
stmsHead Stms SOACS
stms of
    Maybe (Stm SOACS, Stms SOACS)
Nothing -> Stms MC -> ExtractM (Stms MC)
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stms MC
forall a. Monoid a => a
mempty
    Just (Stm SOACS
stm, Stms SOACS
stms') -> do
      stm_stms <- Stm SOACS -> ExtractM (Stms MC)
transformStm Stm SOACS
stm
      inScopeOf stm_stms $ (stm_stms <>) <$> transformStms stms'

transformBody :: Body SOACS -> ExtractM (Body MC)
transformBody :: Body SOACS -> ExtractM (Body MC)
transformBody (Body () Stms SOACS
stms Result
res) =
  BodyDec MC -> Stms MC -> Result -> Body MC
forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
Body () (Stms MC -> Result -> Body MC)
-> ExtractM (Stms MC) -> ExtractM (Result -> Body MC)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stms SOACS -> ExtractM (Stms MC)
transformStms Stms SOACS
stms ExtractM (Result -> Body MC)
-> ExtractM Result -> ExtractM (Body MC)
forall a b. ExtractM (a -> b) -> ExtractM a -> ExtractM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Result -> ExtractM Result
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res

sequentialiseBody :: Body SOACS -> ExtractM (Body MC)
sequentialiseBody :: Body SOACS -> ExtractM (Body MC)
sequentialiseBody = Body MC -> ExtractM (Body MC)
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body MC -> ExtractM (Body MC))
-> (Body SOACS -> Body MC) -> Body SOACS -> ExtractM (Body MC)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Identity (Body MC) -> Body MC
forall a. Identity a -> a
runIdentity (Identity (Body MC) -> Body MC)
-> (Body SOACS -> Identity (Body MC)) -> Body SOACS -> Body MC
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rephraser Identity SOACS MC -> Body SOACS -> Identity (Body MC)
forall (m :: * -> *) from to.
Monad m =>
Rephraser m from to -> Body from -> m (Body to)
rephraseBody Rephraser Identity SOACS MC
toMC
  where
    toMC :: Rephraser Identity SOACS MC
toMC = (SOAC MC -> Op MC) -> Rephraser Identity SOACS MC
forall (m :: * -> *) from to.
(Monad m, SameScope from to, ExpDec from ~ ExpDec to,
 BodyDec from ~ BodyDec to, RetType from ~ RetType to,
 BranchType from ~ BranchType to, Op from ~ SOAC from) =>
(SOAC to -> Op to) -> Rephraser m from to
injectSOACS SOAC MC -> Op MC
SOAC MC -> MCOp SOAC MC
forall (op :: * -> *) rep. op rep -> MCOp op rep
OtherOp

transformFunDef :: FunDef SOACS -> ExtractM (FunDef MC)
transformFunDef :: FunDef SOACS -> ExtractM (FunDef MC)
transformFunDef (FunDef Maybe EntryPoint
entry Attrs
attrs Name
name [(RetType SOACS, RetAls)]
rettype [FParam SOACS]
params Body SOACS
body) = do
  body' <- Scope MC -> ExtractM (Body MC) -> ExtractM (Body MC)
forall a. Scope MC -> ExtractM a -> ExtractM a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([Param (FParamInfo MC)] -> Scope MC
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams [FParam SOACS]
[Param (FParamInfo MC)]
params) (ExtractM (Body MC) -> ExtractM (Body MC))
-> ExtractM (Body MC) -> ExtractM (Body MC)
forall a b. (a -> b) -> a -> b
$ Body SOACS -> ExtractM (Body MC)
transformBody Body SOACS
body
  pure $ FunDef entry attrs name rettype params body'

-- Code generation for each parallel basic block is parameterised over
-- how we handle parallelism in the body (whether it's sequentialised
-- by keeping it as SOACs, or turned into SegOps).

data NeedsRename = DoRename | DoNotRename

renameIfNeeded :: (Rename a) => NeedsRename -> a -> ExtractM a
renameIfNeeded :: forall a. Rename a => NeedsRename -> a -> ExtractM a
renameIfNeeded NeedsRename
DoRename = a -> ExtractM a
forall a (m :: * -> *). (Rename a, MonadFreshNames m) => a -> m a
renameSomething
renameIfNeeded NeedsRename
DoNotRename = a -> ExtractM a
forall a. a -> ExtractM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

transformMap ::
  NeedsRename ->
  (Body SOACS -> ExtractM (Body MC)) ->
  SubExp ->
  Lambda SOACS ->
  [VName] ->
  ExtractM (SegOp () MC)
transformMap :: NeedsRename
-> (Body SOACS -> ExtractM (Body MC))
-> SubExp
-> Lambda SOACS
-> [VName]
-> ExtractM (SegOp () MC)
transformMap NeedsRename
rename Body SOACS -> ExtractM (Body MC)
onBody SubExp
w Lambda SOACS
map_lam [VName]
arrs = do
  (gtid, space) <- SubExp -> ExtractM (VName, SegSpace)
forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> m (VName, SegSpace)
mkSegSpace SubExp
w
  kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs
  renameIfNeeded rename $
    SegMap () space (lambdaReturnType map_lam) kbody

transformRedomap ::
  NeedsRename ->
  (Body SOACS -> ExtractM (Body MC)) ->
  SubExp ->
  [Reduce SOACS] ->
  Lambda SOACS ->
  [VName] ->
  ExtractM ([Stms MC], SegOp () MC)
transformRedomap :: NeedsRename
-> (Body SOACS -> ExtractM (Body MC))
-> SubExp
-> [Reduce SOACS]
-> Lambda SOACS
-> [VName]
-> ExtractM ([Stms MC], SegOp () MC)
transformRedomap NeedsRename
rename Body SOACS -> ExtractM (Body MC)
onBody SubExp
w [Reduce SOACS]
reds Lambda SOACS
map_lam [VName]
arrs = do
  (gtid, space) <- SubExp -> ExtractM (VName, SegSpace)
forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> m (VName, SegSpace)
mkSegSpace SubExp
w
  kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs
  (reds_stms, reds') <- mapAndUnzipM reduceToSegBinOp reds
  op' <-
    renameIfNeeded rename $
      SegRed () space (lambdaReturnType map_lam) kbody reds'
  pure (reds_stms, op')

transformHist ::
  NeedsRename ->
  (Body SOACS -> ExtractM (Body MC)) ->
  SubExp ->
  [SOACS.HistOp SOACS] ->
  Lambda SOACS ->
  [VName] ->
  ExtractM ([Stms MC], SegOp () MC)
transformHist :: NeedsRename
-> (Body SOACS -> ExtractM (Body MC))
-> SubExp
-> [HistOp SOACS]
-> Lambda SOACS
-> [VName]
-> ExtractM ([Stms MC], SegOp () MC)
transformHist NeedsRename
rename Body SOACS -> ExtractM (Body MC)
onBody SubExp
w [HistOp SOACS]
hists Lambda SOACS
map_lam [VName]
arrs = do
  (gtid, space) <- SubExp -> ExtractM (VName, SegSpace)
forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> m (VName, SegSpace)
mkSegSpace SubExp
w
  kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs
  (hists_stms, hists') <- mapAndUnzipM histToSegBinOp hists
  op' <-
    renameIfNeeded rename $
      SegHist () space (lambdaReturnType map_lam) kbody hists'
  pure (hists_stms, op')

transformSOAC :: Pat Type -> Attrs -> SOAC SOACS -> ExtractM (Stms MC)
transformSOAC :: Pat Type -> Attrs -> SOAC SOACS -> ExtractM (Stms MC)
transformSOAC Pat Type
_ Attrs
_ JVP {} =
  String -> ExtractM (Stms MC)
forall a. HasCallStack => String -> a
error String
"transformSOAC: unhandled JVP"
transformSOAC Pat Type
_ Attrs
_ VJP {} =
  String -> ExtractM (Stms MC)
forall a. HasCallStack => String -> a
error String
"transformSOAC: unhandled VJP"
transformSOAC Pat Type
pat Attrs
_ (Screma SubExp
w [VName]
arrs ScremaForm SOACS
form)
  | Just Lambda SOACS
lam <- ScremaForm SOACS -> Maybe (Lambda SOACS)
forall rep. ScremaForm rep -> Maybe (Lambda rep)
isMapSOAC ScremaForm SOACS
form = do
      seq_op <- NeedsRename
-> (Body SOACS -> ExtractM (Body MC))
-> SubExp
-> Lambda SOACS
-> [VName]
-> ExtractM (SegOp () MC)
transformMap NeedsRename
DoNotRename Body SOACS -> ExtractM (Body MC)
sequentialiseBody SubExp
w Lambda SOACS
lam [VName]
arrs
      if lambdaContainsParallelism lam
        then do
          par_op <- transformMap DoRename transformBody w lam arrs
          pure $ oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op)
        else pure $ oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op)
  | Just ([Reduce SOACS]
reds, Lambda SOACS
map_lam) <- ScremaForm SOACS -> Maybe ([Reduce SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
isRedomapSOAC ScremaForm SOACS
form = do
      (seq_reds_stms, seq_op) <-
        NeedsRename
-> (Body SOACS -> ExtractM (Body MC))
-> SubExp
-> [Reduce SOACS]
-> Lambda SOACS
-> [VName]
-> ExtractM ([Stms MC], SegOp () MC)
transformRedomap NeedsRename
DoNotRename Body SOACS -> ExtractM (Body MC)
sequentialiseBody SubExp
w [Reduce SOACS]
reds Lambda SOACS
map_lam [VName]
arrs
      if lambdaContainsParallelism map_lam
        then do
          (par_reds_stms, par_op) <-
            transformRedomap DoRename transformBody w reds map_lam arrs
          pure $
            mconcat (seq_reds_stms <> par_reds_stms)
              <> oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op)
        else
          pure $
            mconcat seq_reds_stms
              <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op)
  | Just ([Scan SOACS]
scans, Lambda SOACS
map_lam) <- ScremaForm SOACS -> Maybe ([Scan SOACS], Lambda SOACS)
forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
isScanomapSOAC ScremaForm SOACS
form = do
      (gtid, space) <- SubExp -> ExtractM (VName, SegSpace)
forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> m (VName, SegSpace)
mkSegSpace SubExp
w
      kbody <- mapLambdaToKernelBody transformBody gtid map_lam arrs
      (scans_stms, scans') <- mapAndUnzipM scanToSegBinOp scans
      pure $
        mconcat scans_stms
          <> oneStm
            ( Let pat (defAux ()) $
                Op $
                  ParOp Nothing $
                    SegScan () space (lambdaReturnType map_lam) kbody scans'
            )
  | Bool
otherwise = do
      -- This screma is too complicated for us to immediately do
      -- anything, so split it up and try again.
      scope <- Scope MC -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope MC -> Scope SOACS)
-> ExtractM (Scope MC) -> ExtractM (Scope SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExtractM (Scope MC)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
      transformStms =<< runBuilderT_ (dissectScrema pat w form arrs) scope
transformSOAC Pat Type
pat Attrs
_ (Scatter SubExp
w [VName]
ivs ScatterSpec VName
dests Lambda SOACS
lam) = do
  (gtid, space) <- SubExp -> ExtractM (VName, SegSpace)
forall (m :: * -> *).
MonadFreshNames m =>
SubExp -> m (VName, SegSpace)
mkSegSpace SubExp
w

  Body () kstms res <- mapLambdaToBody transformBody gtid lam ivs

  (rets, kres) <- fmap unzip $ forM (groupScatterResults dests res) $ \(ShapeBase SubExp
_a_w, VName
a, [(Result, SubExpRes)]
is_vs) -> do
    a_t <- VName -> ExtractM Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
a
    let cs =
          ((Result, SubExpRes) -> Certs) -> [(Result, SubExpRes)] -> Certs
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((SubExpRes -> Certs) -> Result -> Certs
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap SubExpRes -> Certs
resCerts (Result -> Certs)
-> ((Result, SubExpRes) -> Result) -> (Result, SubExpRes) -> Certs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Result, SubExpRes) -> Result
forall a b. (a, b) -> a
fst) [(Result, SubExpRes)]
is_vs
            Certs -> Certs -> Certs
forall a. Semigroup a => a -> a -> a
<> ((Result, SubExpRes) -> Certs) -> [(Result, SubExpRes)] -> Certs
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (SubExpRes -> Certs
resCerts (SubExpRes -> Certs)
-> ((Result, SubExpRes) -> SubExpRes)
-> (Result, SubExpRes)
-> Certs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Result, SubExpRes) -> SubExpRes
forall a b. (a, b) -> b
snd) [(Result, SubExpRes)]
is_vs
        is_vs' = [(Type -> [DimIndex SubExp] -> Slice SubExp
fullSlice Type
a_t ([DimIndex SubExp] -> Slice SubExp)
-> [DimIndex SubExp] -> Slice SubExp
forall a b. (a -> b) -> a -> b
$ (SubExpRes -> DimIndex SubExp) -> Result -> [DimIndex SubExp]
forall a b. (a -> b) -> [a] -> [b]
map (SubExp -> DimIndex SubExp
forall d. d -> DimIndex d
DimFix (SubExp -> DimIndex SubExp)
-> (SubExpRes -> SubExp) -> SubExpRes -> DimIndex SubExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExpRes -> SubExp
resSubExp) Result
is, SubExpRes -> SubExp
resSubExp SubExpRes
v) | (Result
is, SubExpRes
v) <- [(Result, SubExpRes)]
is_vs]
    pure (a_t, WriteReturns cs a is_vs')
  pure . oneStm . Let pat (defAux ()) . Op . ParOp Nothing $
    SegMap () space rets (KernelBody () kstms kres)
transformSOAC Pat Type
pat Attrs
_ (Hist SubExp
w [VName]
arrs [HistOp SOACS]
hists Lambda SOACS
map_lam) = do
  (seq_hist_stms, seq_op) <-
    NeedsRename
-> (Body SOACS -> ExtractM (Body MC))
-> SubExp
-> [HistOp SOACS]
-> Lambda SOACS
-> [VName]
-> ExtractM ([Stms MC], SegOp () MC)
transformHist NeedsRename
DoNotRename Body SOACS -> ExtractM (Body MC)
sequentialiseBody SubExp
w [HistOp SOACS]
hists Lambda SOACS
map_lam [VName]
arrs

  if lambdaContainsParallelism map_lam
    then do
      (par_hist_stms, par_op) <-
        transformHist DoRename transformBody w hists map_lam arrs
      pure $
        mconcat (seq_hist_stms <> par_hist_stms)
          <> oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op)
    else
      pure $
        mconcat seq_hist_stms
          <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op)
transformSOAC Pat Type
pat Attrs
_ (Stream SubExp
w [VName]
arrs [SubExp]
nes Lambda SOACS
lam) = do
  -- Just remove the stream and transform the resulting stms.
  soacs_scope <- Scope MC -> Scope SOACS
forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope (Scope MC -> Scope SOACS)
-> ExtractM (Scope MC) -> ExtractM (Scope SOACS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExtractM (Scope MC)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  stream_stms <-
    flip runBuilderT_ soacs_scope $
      sequentialStreamWholeArray pat w nes lam arrs
  transformStms stream_stms

transformProg :: Prog SOACS -> PassM (Prog MC)
transformProg :: Prog SOACS -> PassM (Prog MC)
transformProg Prog SOACS
prog =
  (VNameSource -> (Prog MC, VNameSource)) -> PassM (Prog MC)
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (Prog MC, VNameSource)) -> PassM (Prog MC))
-> (VNameSource -> (Prog MC, VNameSource)) -> PassM (Prog MC)
forall a b. (a -> b) -> a -> b
$ State VNameSource (Prog MC)
-> VNameSource -> (Prog MC, VNameSource)
forall s a. State s a -> s -> (a, s)
runState (ReaderT (Scope MC) (State VNameSource) (Prog MC)
-> Scope MC -> State VNameSource (Prog MC)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Scope MC) (State VNameSource) (Prog MC)
m Scope MC
forall a. Monoid a => a
mempty)
  where
    ExtractM ReaderT (Scope MC) (State VNameSource) (Prog MC)
m = do
      consts' <- Stms SOACS -> ExtractM (Stms MC)
transformStms (Stms SOACS -> ExtractM (Stms MC))
-> Stms SOACS -> ExtractM (Stms MC)
forall a b. (a -> b) -> a -> b
$ Prog SOACS -> Stms SOACS
forall rep. Prog rep -> Stms rep
progConsts Prog SOACS
prog
      funs' <- inScopeOf consts' $ mapM transformFunDef $ progFuns prog
      pure $
        prog
          { progConsts = consts',
            progFuns = funs'
          }

-- | Transform a program using SOACs to a program in the 'MC'
-- representation, using some amount of flattening.
extractMulticore :: Pass SOACS MC
extractMulticore :: Pass SOACS MC
extractMulticore =
  Pass
    { passName :: String
passName = String
"extract multicore parallelism",
      passDescription :: String
passDescription = String
"Extract multicore parallelism",
      passFunction :: Prog SOACS -> PassM (Prog MC)
passFunction = Prog SOACS -> PassM (Prog MC)
transformProg
    }