{-# LANGUAGE LambdaCase #-}

-- | Add explicit parameters for global names used in device functions.
--
-- The purpose of this pass is to ensure that no functions called from within a
-- parallel section (identified by `SegMap`, `SegRed`, `SegScan, `SegHist`, or
-- `GPUBody`) reference a global name directly. We call these functions "device
-- functions". A global name is a name bound by a statement outside of any
-- defunction definition. This is done by:
--
-- 1. Identifying which global names are referenced by each device function.
-- Since device functions can call each other, this must be determined
-- transitively, using a call graph.
--
-- 2. Adding new parameters to each device function corresponding to each of the
-- global names used by that function.
--
-- 3. Updating each application of a device function to include these new names.
--
-- The main complication is that it is not allowed for a device function to have
-- parameter names that are the same as global names - shadowing is not allowed.
-- However, it is legal for different device functions to have the same
-- parameter names. Hence, once it has been determined which global names are
-- used in device functions, we compute a map that associates the global names
-- to parameter names, and use these. This implies that these names must also be
-- substituted into the device function bodies. Note that when modifying the
-- application of the outermost device function call with additional function
-- arguments, we must use the original global names.
module Futhark.Pass.AddGlobalParams (addGlobalParams) where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.State.Strict
import Data.Map.Strict qualified as M
import Data.Maybe (mapMaybe)
import Data.Set qualified as S
import Futhark.IR.GPU
import Futhark.MonadFreshNames
import Futhark.Pass
import Futhark.Transform.Substitute

data CallMode = AllCalls | ParallelCalls
  deriving (CallMode -> CallMode -> Bool
(CallMode -> CallMode -> Bool)
-> (CallMode -> CallMode -> Bool) -> Eq CallMode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CallMode -> CallMode -> Bool
== :: CallMode -> CallMode -> Bool
$c/= :: CallMode -> CallMode -> Bool
/= :: CallMode -> CallMode -> Bool
Eq)

callsInProg :: CallMode -> Prog GPU -> S.Set Name
callsInProg :: CallMode -> Prog GPU -> Set Name
callsInProg CallMode
mode Prog GPU
prog =
  CallMode -> Bool -> Stms GPU -> Set Name
callsInStms CallMode
mode Bool
False (Prog GPU -> Stms GPU
forall rep. Prog rep -> Stms rep
progConsts Prog GPU
prog)
    Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> (FunDef GPU -> Set Name) -> [FunDef GPU] -> Set Name
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (CallMode -> Bool -> Body GPU -> Set Name
callsInBody CallMode
mode Bool
False (Body GPU -> Set Name)
-> (FunDef GPU -> Body GPU) -> FunDef GPU -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef GPU -> Body GPU
forall rep. FunDef rep -> Body rep
funDefBody) (Prog GPU -> [FunDef GPU]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog GPU
prog)

calledInParallel :: Prog GPU -> S.Set Name
calledInParallel :: Prog GPU -> Set Name
calledInParallel = CallMode -> Prog GPU -> Set Name
callsInProg CallMode
ParallelCalls

callsInGBody :: CallMode -> Bool -> GBody GPU res -> S.Set Name
callsInGBody :: forall res. CallMode -> Bool -> GBody GPU res -> Set Name
callsInGBody CallMode
mode Bool
in_parallel = CallMode -> Bool -> Stms GPU -> Set Name
callsInStms CallMode
mode Bool
in_parallel (Stms GPU -> Set Name)
-> (GBody GPU res -> Stms GPU) -> GBody GPU res -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GBody GPU res -> Stms GPU
forall rep res. GBody rep res -> Stms rep
bodyStms

callsInBody :: CallMode -> Bool -> Body GPU -> S.Set Name
callsInBody :: CallMode -> Bool -> Body GPU -> Set Name
callsInBody = CallMode -> Bool -> Body GPU -> Set Name
forall res. CallMode -> Bool -> GBody GPU res -> Set Name
callsInGBody

callsInKernelBody :: CallMode -> KernelBody GPU -> S.Set Name
callsInKernelBody :: CallMode -> KernelBody GPU -> Set Name
callsInKernelBody CallMode
mode = CallMode -> Bool -> KernelBody GPU -> Set Name
forall res. CallMode -> Bool -> GBody GPU res -> Set Name
callsInGBody CallMode
mode Bool
True

callsInStms :: CallMode -> Bool -> Stms GPU -> S.Set Name
callsInStms :: CallMode -> Bool -> Stms GPU -> Set Name
callsInStms CallMode
mode Bool
in_parallel =
  (Stm GPU -> Set Name) -> [Stm GPU] -> Set Name
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (CallMode -> Bool -> Exp GPU -> Set Name
callsInExp CallMode
mode Bool
in_parallel (Exp GPU -> Set Name)
-> (Stm GPU -> Exp GPU) -> Stm GPU -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm GPU -> Exp GPU
forall rep. Stm rep -> Exp rep
stmExp) ([Stm GPU] -> Set Name)
-> (Stms GPU -> [Stm GPU]) -> Stms GPU -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms GPU -> [Stm GPU]
forall rep. Stms rep -> [Stm rep]
stmsToList

callsInExp :: CallMode -> Bool -> Exp GPU -> S.Set Name
callsInExp :: CallMode -> Bool -> Exp GPU -> Set Name
callsInExp CallMode
mode Bool
in_parallel = \case
  Apply Name
fname [(SubExp, Diet)]
_ [(RetType GPU, RetAls)]
_ Safety
_
    | CallMode
mode CallMode -> CallMode -> Bool
forall a. Eq a => a -> a -> Bool
== CallMode
AllCalls Bool -> Bool -> Bool
|| Bool
in_parallel -> Name -> Set Name
forall a. a -> Set a
S.singleton Name
fname
    | Bool
otherwise -> Set Name
forall a. Monoid a => a
mempty
  Match [SubExp]
_ [Case (Body GPU)]
cases Body GPU
defbody MatchDec (BranchType GPU)
_ ->
    (Case (Body GPU) -> Set Name) -> [Case (Body GPU)] -> Set Name
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (CallMode -> Bool -> Body GPU -> Set Name
callsInBody CallMode
mode Bool
in_parallel (Body GPU -> Set Name)
-> (Case (Body GPU) -> Body GPU) -> Case (Body GPU) -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Case (Body GPU) -> Body GPU
forall body. Case body -> body
caseBody) [Case (Body GPU)]
cases
      Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> CallMode -> Bool -> Body GPU -> Set Name
callsInBody CallMode
mode Bool
in_parallel Body GPU
defbody
  Loop [(FParam GPU, SubExp)]
_ LoopForm
_ Body GPU
body ->
    CallMode -> Bool -> Body GPU -> Set Name
callsInBody CallMode
mode Bool
in_parallel Body GPU
body
  WithAcc [WithAccInput GPU]
inputs Lambda GPU
lam ->
    (WithAccInput GPU -> Set Name) -> [WithAccInput GPU] -> Set Name
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap
      ( \(Shape
_, [VName]
_, Maybe (Lambda GPU, [SubExp])
op) ->
          Set Name
-> ((Lambda GPU, [SubExp]) -> Set Name)
-> Maybe (Lambda GPU, [SubExp])
-> Set Name
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set Name
forall a. Monoid a => a
mempty (\(Lambda GPU
f, [SubExp]
_) -> CallMode -> Bool -> Lambda GPU -> Set Name
callsInLambda CallMode
mode Bool
in_parallel Lambda GPU
f) Maybe (Lambda GPU, [SubExp])
op
      )
      [WithAccInput GPU]
inputs
      Set Name -> Set Name -> Set Name
forall a. Semigroup a => a -> a -> a
<> CallMode -> Bool -> Lambda GPU -> Set Name
callsInLambda CallMode
mode Bool
in_parallel Lambda GPU
lam
  Op Op GPU
op ->
    CallMode -> Bool -> Op GPU -> Set Name
callsInOp CallMode
mode Bool
in_parallel Op GPU
op
  Exp GPU
_ ->
    Set Name
forall a. Monoid a => a
mempty

callsInLambda :: CallMode -> Bool -> Lambda GPU -> S.Set Name
callsInLambda :: CallMode -> Bool -> Lambda GPU -> Set Name
callsInLambda CallMode
mode Bool
in_parallel (Lambda [LParam GPU]
_ [Type]
_ Body GPU
body) = CallMode -> Bool -> Body GPU -> Set Name
callsInBody CallMode
mode Bool
in_parallel Body GPU
body

tell :: (MonadState s m, Semigroup s) => s -> m ()
tell :: forall s (m :: * -> *). (MonadState s m, Semigroup s) => s -> m ()
tell s
x = (s -> s) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (s -> s -> s
forall a. Semigroup a => a -> a -> a
<> s
x)

callsInSOAC :: CallMode -> Bool -> SOAC GPU -> S.Set Name
callsInSOAC :: CallMode -> Bool -> SOAC GPU -> Set Name
callsInSOAC CallMode
mode Bool
in_parallel SOAC GPU
soac =
  State (Set Name) () -> Set Name -> Set Name
forall s a. State s a -> s -> s
execState (StateT (Set Name) Identity (SOAC GPU) -> State (Set Name) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set Name) Identity (SOAC GPU) -> State (Set Name) ())
-> StateT (Set Name) Identity (SOAC GPU) -> State (Set Name) ()
forall a b. (a -> b) -> a -> b
$ SOACMapper GPU GPU (StateT (Set Name) Identity)
-> SOAC GPU -> StateT (Set Name) Identity (SOAC GPU)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper GPU GPU (StateT (Set Name) Identity)
mapper SOAC GPU
soac) Set Name
forall a. Monoid a => a
mempty
  where
    mapper :: SOACMapper GPU GPU (StateT (Set Name) Identity)
mapper =
      SOACMapper GPU GPU (StateT (Set Name) Identity)
forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper
        { mapOnSOACLambda = \Lambda GPU
lam -> do
            Set Name -> State (Set Name) ()
forall s (m :: * -> *). (MonadState s m, Semigroup s) => s -> m ()
tell (Set Name -> State (Set Name) ())
-> Set Name -> State (Set Name) ()
forall a b. (a -> b) -> a -> b
$ CallMode -> Bool -> Lambda GPU -> Set Name
callsInLambda CallMode
mode Bool
in_parallel Lambda GPU
lam
            Lambda GPU -> StateT (Set Name) Identity (Lambda GPU)
forall a. a -> StateT (Set Name) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda GPU
lam
        }

callsInSegOp :: CallMode -> SegOp SegLevel GPU -> S.Set Name
callsInSegOp :: CallMode -> SegOp SegLevel GPU -> Set Name
callsInSegOp CallMode
mode SegOp SegLevel GPU
segop =
  State (Set Name) () -> Set Name -> Set Name
forall s a. State s a -> s -> s
execState (StateT (Set Name) Identity (SegOp SegLevel GPU)
-> State (Set Name) ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (StateT (Set Name) Identity (SegOp SegLevel GPU)
 -> State (Set Name) ())
-> StateT (Set Name) Identity (SegOp SegLevel GPU)
-> State (Set Name) ()
forall a b. (a -> b) -> a -> b
$ SegOpMapper SegLevel GPU GPU (StateT (Set Name) Identity)
-> SegOp SegLevel GPU
-> StateT (Set Name) Identity (SegOp SegLevel GPU)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper SegLevel GPU GPU (StateT (Set Name) Identity)
forall {lvl}. SegOpMapper lvl GPU GPU (StateT (Set Name) Identity)
mapper SegOp SegLevel GPU
segop) Set Name
forall a. Monoid a => a
mempty
  where
    mapper :: SegOpMapper lvl GPU GPU (StateT (Set Name) Identity)
mapper =
      SegOpMapper
  lvl (ZonkAny 0) (ZonkAny 0) (StateT (Set Name) Identity)
forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
identitySegOpMapper
        { mapOnSegBinOpLambda = \Lambda GPU
lam -> do
            Set Name -> State (Set Name) ()
forall s (m :: * -> *). (MonadState s m, Semigroup s) => s -> m ()
tell (Set Name -> State (Set Name) ())
-> Set Name -> State (Set Name) ()
forall a b. (a -> b) -> a -> b
$ CallMode -> Bool -> Lambda GPU -> Set Name
callsInLambda CallMode
mode Bool
True Lambda GPU
lam
            Lambda GPU -> StateT (Set Name) Identity (Lambda GPU)
forall a. a -> StateT (Set Name) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda GPU
lam,
          mapOnSegPostOpLambda = \Lambda GPU
lam -> do
            Set Name -> State (Set Name) ()
forall s (m :: * -> *). (MonadState s m, Semigroup s) => s -> m ()
tell (Set Name -> State (Set Name) ())
-> Set Name -> State (Set Name) ()
forall a b. (a -> b) -> a -> b
$ CallMode -> Bool -> Lambda GPU -> Set Name
callsInLambda CallMode
mode Bool
True Lambda GPU
lam
            Lambda GPU -> StateT (Set Name) Identity (Lambda GPU)
forall a. a -> StateT (Set Name) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Lambda GPU
lam,
          mapOnSegOpBody = \KernelBody GPU
body -> do
            Set Name -> State (Set Name) ()
forall s (m :: * -> *). (MonadState s m, Semigroup s) => s -> m ()
tell (Set Name -> State (Set Name) ())
-> Set Name -> State (Set Name) ()
forall a b. (a -> b) -> a -> b
$ CallMode -> KernelBody GPU -> Set Name
callsInKernelBody CallMode
mode KernelBody GPU
body
            KernelBody GPU -> StateT (Set Name) Identity (KernelBody GPU)
forall a. a -> StateT (Set Name) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure KernelBody GPU
body
        }

callsInOp :: CallMode -> Bool -> Op GPU -> S.Set Name
callsInOp :: CallMode -> Bool -> Op GPU -> Set Name
callsInOp CallMode
mode Bool
in_parallel = \case
  SegOp SegOp SegLevel GPU
segop ->
    CallMode -> SegOp SegLevel GPU -> Set Name
callsInSegOp CallMode
mode SegOp SegLevel GPU
segop
  OtherOp SOAC GPU
soac ->
    CallMode -> Bool -> SOAC GPU -> Set Name
callsInSOAC CallMode
mode Bool
in_parallel SOAC GPU
soac
  GPUBody [Type]
_ Body GPU
body ->
    CallMode -> Bool -> Body GPU -> Set Name
callsInBody CallMode
mode Bool
True Body GPU
body
  Op GPU
_ ->
    Set Name
forall a. Monoid a => a
mempty

buildCallGraphGPU :: Prog GPU -> M.Map Name (S.Set Name)
buildCallGraphGPU :: Prog GPU -> Map Name (Set Name)
buildCallGraphGPU =
  [(Name, Set Name)] -> Map Name (Set Name)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
    ([(Name, Set Name)] -> Map Name (Set Name))
-> (Prog GPU -> [(Name, Set Name)])
-> Prog GPU
-> Map Name (Set Name)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (FunDef GPU -> (Name, Set Name))
-> [FunDef GPU] -> [(Name, Set Name)]
forall a b. (a -> b) -> [a] -> [b]
map (\FunDef GPU
fd -> (FunDef GPU -> Name
forall rep. FunDef rep -> Name
funDefName FunDef GPU
fd, CallMode -> Bool -> Body GPU -> Set Name
callsInBody CallMode
AllCalls Bool
False (Body GPU -> Set Name) -> Body GPU -> Set Name
forall a b. (a -> b) -> a -> b
$ FunDef GPU -> Body GPU
forall rep. FunDef rep -> Body rep
funDefBody FunDef GPU
fd))
    ([FunDef GPU] -> [(Name, Set Name)])
-> (Prog GPU -> [FunDef GPU]) -> Prog GPU -> [(Name, Set Name)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog GPU -> [FunDef GPU]
forall rep. Prog rep -> [FunDef rep]
progFuns

transitiveClosure :: (Ord k) => M.Map k (S.Set k) -> S.Set k -> S.Set k
transitiveClosure :: forall k. Ord k => Map k (Set k) -> Set k -> Set k
transitiveClosure Map k (Set k)
graph = Set k -> Set k
go
  where
    go :: Set k -> Set k
go Set k
seen =
      let seen' :: Set k
seen' =
            Set k
seen
              Set k -> Set k -> Set k
forall a. Semigroup a => a -> a -> a
<> [Set k] -> Set k
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions
                [Set k -> k -> Map k (Set k) -> Set k
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Set k
forall a. Monoid a => a
mempty k
f Map k (Set k)
graph | k
f <- Set k -> [k]
forall a. Set a -> [a]
S.toList Set k
seen]
       in if Set k
seen' Set k -> Set k -> Bool
forall a. Eq a => a -> a -> Bool
== Set k
seen then Set k
seen else Set k -> Set k
go Set k
seen'

globalsPerFun ::
  M.Map Name (S.Set Name) ->
  M.Map Name (S.Set VName) ->
  M.Map Name (S.Set VName)
globalsPerFun :: Map Name (Set Name) -> Map Name (Set VName) -> Map Name (Set VName)
globalsPerFun Map Name (Set Name)
call_graph = Map Name (Set VName) -> Map Name (Set VName)
forall {a}. Ord a => Map Name (Set a) -> Map Name (Set a)
fixpoint
  where
    fixpoint :: Map Name (Set a) -> Map Name (Set a)
fixpoint Map Name (Set a)
m =
      let step :: Name -> Set a -> Set a
step Name
f Set a
gs =
            Set a
gs
              Set a -> Set a -> Set a
forall a. Semigroup a => a -> a -> a
<> [Set a] -> Set a
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions
                [ Set a -> Name -> Map Name (Set a) -> Set a
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Set a
forall a. Monoid a => a
mempty Name
g Map Name (Set a)
m
                | Name
g <- Set Name -> [Name]
forall a. Set a -> [a]
S.toList (Set Name -> [Name]) -> Set Name -> [Name]
forall a b. (a -> b) -> a -> b
$ Set Name -> Name -> Map Name (Set Name) -> Set Name
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Set Name
forall a. Monoid a => a
mempty Name
f Map Name (Set Name)
call_graph
                ]
          m' :: Map Name (Set a)
m' = (Name -> Set a -> Set a) -> Map Name (Set a) -> Map Name (Set a)
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey Name -> Set a -> Set a
step Map Name (Set a)
m
       in if Map Name (Set a)
m' Map Name (Set a) -> Map Name (Set a) -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name (Set a)
m then Map Name (Set a)
m else Map Name (Set a) -> Map Name (Set a)
fixpoint Map Name (Set a)
m'

globalTypes :: Stms GPU -> M.Map VName DeclType
globalTypes :: Stms GPU -> Map VName DeclType
globalTypes =
  [(VName, DeclType)] -> Map VName DeclType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
    ([(VName, DeclType)] -> Map VName DeclType)
-> (Stms GPU -> [(VName, DeclType)])
-> Stms GPU
-> Map VName DeclType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm GPU -> [(VName, DeclType)])
-> [Stm GPU] -> [(VName, DeclType)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap
      ( (PatElem Type -> (VName, DeclType))
-> [PatElem Type] -> [(VName, DeclType)]
forall a b. (a -> b) -> [a] -> [b]
map
          (\PatElem Type
pe -> (PatElem Type -> VName
forall dec. PatElem dec -> VName
patElemName PatElem Type
pe, Type -> Diet -> DeclType
forall shape. TypeBase shape NoMode -> Diet -> TypeBase shape Diet
toDecl (PatElem Type -> Type
forall dec. Typed dec => PatElem dec -> Type
patElemType PatElem Type
pe) Diet
Observe))
          ([PatElem Type] -> [(VName, DeclType)])
-> (Stm GPU -> [PatElem Type]) -> Stm GPU -> [(VName, DeclType)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat Type -> [PatElem Type]
forall dec. Pat dec -> [PatElem dec]
patElems
          (Pat Type -> [PatElem Type])
-> (Stm GPU -> Pat Type) -> Stm GPU -> [PatElem Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm GPU -> Pat Type
Stm GPU -> Pat (LetDec GPU)
forall rep. Stm rep -> Pat (LetDec rep)
stmPat
      )
    ([Stm GPU] -> [(VName, DeclType)])
-> (Stms GPU -> [Stm GPU]) -> Stms GPU -> [(VName, DeclType)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms GPU -> [Stm GPU]
forall rep. Stms rep -> [Stm rep]
stmsToList

transformProg :: Prog GPU -> PassM (Prog GPU)
transformProg :: Prog GPU -> PassM (Prog GPU)
transformProg Prog GPU
prog = do
  let global_tps :: Map VName DeclType
global_tps = Stms GPU -> Map VName DeclType
globalTypes (Stms GPU -> Map VName DeclType) -> Stms GPU -> Map VName DeclType
forall a b. (a -> b) -> a -> b
$ Prog GPU -> Stms GPU
forall rep. Prog rep -> Stms rep
progConsts Prog GPU
prog
      globals :: Set VName
globals = Map VName DeclType -> Set VName
forall k a. Map k a -> Set k
M.keysSet Map VName DeclType
global_tps
      call_graph :: Map Name (Set Name)
call_graph = Prog GPU -> Map Name (Set Name)
buildCallGraphGPU Prog GPU
prog
      roots :: Set Name
roots = Prog GPU -> Set Name
calledInParallel Prog GPU
prog
      device_funs :: Set Name
device_funs = Map Name (Set Name) -> Set Name -> Set Name
forall k. Ord k => Map k (Set k) -> Set k -> Set k
transitiveClosure Map Name (Set Name)
call_graph Set Name
roots
      direct_globals :: Map Name (Set VName)
direct_globals =
        [(Name, Set VName)] -> Map Name (Set VName)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
          [ (FunDef GPU -> Name
forall rep. FunDef rep -> Name
funDefName FunDef GPU
fd, [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList (Names -> [VName]
namesToList (Names -> [VName]) -> Names -> [VName]
forall a b. (a -> b) -> a -> b
$ FunDef GPU -> Names
forall a. FreeIn a => a -> Names
freeIn FunDef GPU
fd) Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Set VName
globals)
          | FunDef GPU
fd <- Prog GPU -> [FunDef GPU]
forall rep. Prog rep -> [FunDef rep]
progFuns Prog GPU
prog,
            FunDef GPU -> Name
forall rep. FunDef rep -> Name
funDefName FunDef GPU
fd Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Name
device_funs
          ]
      transitive_globals :: Map Name (Set VName)
transitive_globals = Map Name (Set Name) -> Map Name (Set VName) -> Map Name (Set VName)
globalsPerFun Map Name (Set Name)
call_graph Map Name (Set VName)
direct_globals
      all_used_globals :: Set VName
all_used_globals = [Set VName] -> Set VName
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
S.unions ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ Map Name (Set VName) -> [Set VName]
forall k a. Map k a -> [a]
M.elems Map Name (Set VName)
transitive_globals

  global_to_param <-
    [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, VName)] -> Map VName VName)
-> PassM [(VName, VName)] -> PassM (Map VName VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> PassM (VName, VName))
-> [VName] -> PassM [(VName, VName)]
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 (\VName
g -> (VName
g,) (VName -> (VName, VName)) -> PassM VName -> PassM (VName, VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> PassM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
g) (Set VName -> [VName]
forall a. Set a -> [a]
S.toAscList Set VName
all_used_globals)

  let globals_ordered =
        (Set VName -> [VName]) -> Map Name (Set VName) -> Map Name [VName]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Set VName -> [VName]
forall a. Set a -> [a]
S.toAscList Map Name (Set VName)
transitive_globals
      globals_for_fun Name
f =
        [VName] -> Name -> Map Name [VName] -> [VName]
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault [VName]
forall a. Monoid a => a
mempty Name
f Map Name [VName]
globals_ordered
      params_for =
        (VName -> Maybe (Param DeclType)) -> [VName] -> [Param DeclType]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((VName -> Maybe (Param DeclType)) -> [VName] -> [Param DeclType])
-> (VName -> Maybe (Param DeclType)) -> [VName] -> [Param DeclType]
forall a b. (a -> b) -> a -> b
$ \VName
g -> do
          p <- VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
g Map VName VName
global_to_param
          t <- M.lookup g global_tps
          pure $ Param mempty p t
      use_name Map k k
env k
g = k -> k -> Map k k -> k
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault k
g k
g Map k k
env
      call_args Map VName VName
env Name
f =
        [ (VName -> SubExp
Var (VName -> SubExp) -> VName -> SubExp
forall a b. (a -> b) -> a -> b
$ Map VName VName -> VName -> VName
forall {k}. Ord k => Map k k -> k -> k
use_name Map VName VName
env VName
g, Diet
Observe)
        | VName
g <- Name -> [VName]
globals_for_fun Name
f
        ]
      call_rewriter :: M.Map VName VName -> Exp GPU -> Exp GPU
      call_rewriter Map VName VName
env (Apply Name
fname [(SubExp, Diet)]
args [(RetType GPU, RetAls)]
rettype Safety
safety) =
        Name
-> [(SubExp, Diet)] -> [(RetType GPU, RetAls)] -> Safety -> Exp GPU
forall rep.
Name
-> [(SubExp, Diet)] -> [(RetType rep, RetAls)] -> Safety -> Exp rep
Apply Name
fname ([(SubExp, Diet)]
args [(SubExp, Diet)] -> [(SubExp, Diet)] -> [(SubExp, Diet)]
forall a. Semigroup a => a -> a -> a
<> Map VName VName -> Name -> [(SubExp, Diet)]
call_args Map VName VName
env Name
fname) [(RetType GPU, RetAls)]
rettype Safety
safety
      call_rewriter Map VName VName
env Exp GPU
e = Mapper GPU GPU Identity -> Exp GPU -> Exp GPU
forall frep trep. Mapper frep trep Identity -> Exp frep -> Exp trep
mapExp Mapper GPU GPU Identity
mapper Exp GPU
e
        where
          mapper :: Mapper GPU GPU Identity
          mapper :: Mapper GPU GPU Identity
mapper =
            Mapper
              { mapOnSubExp :: SubExp -> Identity SubExp
mapOnSubExp = SubExp -> Identity SubExp
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
                mapOnBody :: Scope GPU -> Body GPU -> Identity (Body GPU)
mapOnBody = (Body GPU -> Identity (Body GPU))
-> Scope GPU -> Body GPU -> Identity (Body GPU)
forall a b. a -> b -> a
const ((Body GPU -> Identity (Body GPU))
 -> Scope GPU -> Body GPU -> Identity (Body GPU))
-> (Body GPU -> Identity (Body GPU))
-> Scope GPU
-> Body GPU
-> Identity (Body GPU)
forall a b. (a -> b) -> a -> b
$ Body GPU -> Identity (Body GPU)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body GPU -> Identity (Body GPU))
-> (Body GPU -> Body GPU) -> Body GPU -> Identity (Body GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName VName -> Body GPU -> Body GPU
rewriteBody Map VName VName
env,
                mapOnVName :: VName -> Identity VName
mapOnVName = VName -> Identity VName
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
                mapOnRetType :: RetType GPU -> Identity (RetType GPU)
mapOnRetType = ExtType -> Identity ExtType
RetType GPU -> Identity (RetType GPU)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
                mapOnBranchType :: BranchType GPU -> Identity (BranchType GPU)
mapOnBranchType = ExtType -> Identity ExtType
BranchType GPU -> Identity (BranchType GPU)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
                mapOnFParam :: FParam GPU -> Identity (FParam GPU)
mapOnFParam = Param DeclType -> Identity (Param DeclType)
FParam GPU -> Identity (FParam GPU)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
                mapOnLParam :: LParam GPU -> Identity (LParam GPU)
mapOnLParam = Param Type -> Identity (Param Type)
LParam GPU -> Identity (LParam GPU)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
                mapOnOp :: Op GPU -> Identity (Op GPU)
mapOnOp = HostOp SOAC GPU -> Identity (HostOp SOAC GPU)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HostOp SOAC GPU -> Identity (HostOp SOAC GPU))
-> (HostOp SOAC GPU -> HostOp SOAC GPU)
-> HostOp SOAC GPU
-> Identity (HostOp SOAC GPU)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName VName -> Op GPU -> Op GPU
rewriteOp Map VName VName
env
              }
      rewriteStm :: M.Map VName VName -> Stm GPU -> Stm GPU
      rewriteStm Map VName VName
env (Let Pat (LetDec GPU)
pat StmAux (ExpDec GPU)
aux Exp GPU
e) = Pat (LetDec GPU) -> StmAux (ExpDec GPU) -> Exp GPU -> Stm GPU
forall rep.
Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
Let Pat (LetDec GPU)
pat StmAux (ExpDec GPU)
aux (Exp GPU -> Stm GPU) -> Exp GPU -> Stm GPU
forall a b. (a -> b) -> a -> b
$ Map VName VName -> Exp GPU -> Exp GPU
call_rewriter Map VName VName
env Exp GPU
e
      rewriteStms :: M.Map VName VName -> Stms GPU -> Stms GPU
      rewriteStms Map VName VName
env = [Stm GPU] -> Stms GPU
forall rep. [Stm rep] -> Stms rep
stmsFromList ([Stm GPU] -> Stms GPU)
-> (Stms GPU -> [Stm GPU]) -> Stms GPU -> Stms GPU
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm GPU -> Stm GPU) -> [Stm GPU] -> [Stm GPU]
forall a b. (a -> b) -> [a] -> [b]
map (Map VName VName -> Stm GPU -> Stm GPU
rewriteStm Map VName VName
env) ([Stm GPU] -> [Stm GPU])
-> (Stms GPU -> [Stm GPU]) -> Stms GPU -> [Stm GPU]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms GPU -> [Stm GPU]
forall rep. Stms rep -> [Stm rep]
stmsToList
      rewriteGBody :: M.Map VName VName -> GBody GPU res -> GBody GPU res
      rewriteGBody Map VName VName
env (Body BodyDec GPU
dec Stms GPU
stms [res]
res) =
        BodyDec GPU -> Stms GPU -> [res] -> GBody GPU res
forall rep res. BodyDec rep -> Stms rep -> [res] -> GBody rep res
Body BodyDec GPU
dec (Map VName VName -> Stms GPU -> Stms GPU
rewriteStms Map VName VName
env Stms GPU
stms) [res]
res
      rewriteBody :: M.Map VName VName -> Body GPU -> Body GPU
      rewriteBody = Map VName VName -> Body GPU -> Body GPU
forall res. Map VName VName -> GBody GPU res -> GBody GPU res
rewriteGBody
      rewriteKernelBody :: M.Map VName VName -> KernelBody GPU -> KernelBody GPU
      rewriteKernelBody = Map VName VName -> KernelBody GPU -> KernelBody GPU
forall res. Map VName VName -> GBody GPU res -> GBody GPU res
rewriteGBody
      rewriteLambda :: M.Map VName VName -> Lambda GPU -> Lambda GPU
      rewriteLambda Map VName VName
env (Lambda [LParam GPU]
ps [Type]
ret Body GPU
body) =
        [LParam GPU] -> [Type] -> Body GPU -> Lambda GPU
forall rep t. [LParam rep] -> [t] -> Body rep -> GLambda rep t
Lambda [LParam GPU]
ps [Type]
ret (Body GPU -> Lambda GPU) -> Body GPU -> Lambda GPU
forall a b. (a -> b) -> a -> b
$ Map VName VName -> Body GPU -> Body GPU
rewriteBody Map VName VName
env Body GPU
body
      rewriteOp :: M.Map VName VName -> Op GPU -> Op GPU
      rewriteOp Map VName VName
env (SegOp SegOp SegLevel GPU
segop) =
        let segmapper :: SegOpMapper lvl GPU GPU Identity
segmapper =
              SegOpMapper lvl (ZonkAny 1) (ZonkAny 1) Identity
forall (m :: * -> *) lvl rep. Monad m => SegOpMapper lvl rep rep m
identitySegOpMapper
                { mapOnSegBinOpLambda = pure . rewriteLambda env,
                  mapOnSegPostOpLambda = pure . rewriteLambda env,
                  mapOnSegOpBody = pure . rewriteKernelBody env
                }
         in SegOp SegLevel GPU -> HostOp SOAC GPU
forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
SegOp (SegOp SegLevel GPU -> HostOp SOAC GPU)
-> SegOp SegLevel GPU -> HostOp SOAC GPU
forall a b. (a -> b) -> a -> b
$ Identity (SegOp SegLevel GPU) -> SegOp SegLevel GPU
forall a. Identity a -> a
runIdentity (Identity (SegOp SegLevel GPU) -> SegOp SegLevel GPU)
-> Identity (SegOp SegLevel GPU) -> SegOp SegLevel GPU
forall a b. (a -> b) -> a -> b
$ SegOpMapper SegLevel GPU GPU Identity
-> SegOp SegLevel GPU -> Identity (SegOp SegLevel GPU)
forall (m :: * -> *) lvl frep trep.
Monad m =>
SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
mapSegOpM SegOpMapper SegLevel GPU GPU Identity
forall {lvl}. SegOpMapper lvl GPU GPU Identity
segmapper SegOp SegLevel GPU
segop
      rewriteOp Map VName VName
env (OtherOp SOAC GPU
soac) =
        let soacmapper :: SOACMapper GPU GPU Identity
soacmapper =
              SOACMapper GPU GPU Identity
forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
identitySOACMapper
                { mapOnSOACLambda = pure . rewriteLambda env
                }
         in SOAC GPU -> HostOp SOAC GPU
forall (op :: * -> *) rep. op rep -> HostOp op rep
OtherOp (SOAC GPU -> HostOp SOAC GPU) -> SOAC GPU -> HostOp SOAC GPU
forall a b. (a -> b) -> a -> b
$ Identity (SOAC GPU) -> SOAC GPU
forall a. Identity a -> a
runIdentity (Identity (SOAC GPU) -> SOAC GPU)
-> Identity (SOAC GPU) -> SOAC GPU
forall a b. (a -> b) -> a -> b
$ SOACMapper GPU GPU Identity -> SOAC GPU -> Identity (SOAC GPU)
forall (m :: * -> *) frep trep.
Monad m =>
SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
mapSOACM SOACMapper GPU GPU Identity
soacmapper SOAC GPU
soac
      rewriteOp Map VName VName
env (GPUBody [Type]
ts Body GPU
body) =
        [Type] -> Body GPU -> HostOp SOAC GPU
forall (op :: * -> *) rep. [Type] -> Body rep -> HostOp op rep
GPUBody [Type]
ts (Body GPU -> HostOp SOAC GPU) -> Body GPU -> HostOp SOAC GPU
forall a b. (a -> b) -> a -> b
$ Map VName VName -> Body GPU -> Body GPU
rewriteBody Map VName VName
env Body GPU
body
      rewriteOp Map VName VName
_ Op GPU
op = Op GPU
op

      rewriteFun :: FunDef GPU -> FunDef GPU
      rewriteFun FunDef GPU
fd
        | FunDef GPU -> Name
forall rep. FunDef rep -> Name
funDefName FunDef GPU
fd Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set Name
device_funs =
            let gs :: [VName]
gs = Name -> [VName]
globals_for_fun (FunDef GPU -> Name
forall rep. FunDef rep -> Name
funDefName FunDef GPU
fd)
                env :: Map VName VName
env =
                  [(VName, VName)] -> Map VName VName
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                    [ (VName
g, VName
p)
                    | VName
g <- [VName]
gs,
                      Just VName
p <- [VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
g Map VName VName
global_to_param]
                    ]
                substs :: Map VName VName
substs = Map VName VName
env
                extra_params :: [Param DeclType]
extra_params = [VName] -> [Param DeclType]
params_for [VName]
gs
                body' :: Body GPU
body' = Map VName VName -> Body GPU -> Body GPU
forall a. Substitute a => Map VName VName -> a -> a
substituteNames Map VName VName
substs (Body GPU -> Body GPU) -> Body GPU -> Body GPU
forall a b. (a -> b) -> a -> b
$ FunDef GPU -> Body GPU
forall rep. FunDef rep -> Body rep
funDefBody FunDef GPU
fd
             in FunDef GPU
fd
                  { funDefParams = funDefParams fd <> extra_params,
                    funDefBody = rewriteBody env body'
                  }
        | Bool
otherwise =
            FunDef GPU
fd {funDefBody = rewriteBody mempty (funDefBody fd)}

  pure
    prog
      { progConsts = rewriteStms mempty $ progConsts prog,
        progFuns = map rewriteFun $ progFuns prog
      }

-- | Ensure that device functions do not reference global names directly.
addGlobalParams :: Pass GPU GPU
addGlobalParams :: Pass GPU GPU
addGlobalParams =
  Pass
    { passName :: String
passName = String
"add global params",
      passDescription :: String
passDescription = String
"Thread global names explicitly into device functions.",
      passFunction :: Prog GPU -> PassM (Prog GPU)
passFunction = Prog GPU -> PassM (Prog GPU)
transformProg
    }