    -- |
    -- This module implements program analysis to determine which program statements
    -- the "Futhark.Optimise.ReduceDeviceSyncs" pass should move into 'GPUBody' kernels
    -- to reduce blocking memory transfers between host and device. The results of
    -- the analysis is encoded into a 'MigrationTable' which can be queried.
    --
    -- To reduce blocking scalar reads the module constructs a data flow
    -- dependency graph of program variables (see
    -- "Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph") in which
    -- it finds a minimum vertex cut that separates array reads of scalars
    -- from transitive usage that cannot or should not be migrated to
    -- device.
    --
    -- The variables of each partition are assigned a 'MigrationStatus' that states
    -- whether the computation of those variables should be moved to device or
    -- remain on host. Due to how the graph is built and the vertex cut is found all
    -- variables bound by a single statement will belong to the same partition.
    --
    -- The vertex cut contains all variables that will reside in device memory but
    -- are required by host operations. These variables must be read from device
    -- memory and cannot be reduced further in number merely by migrating
    -- statements (subject to the accuracy of the graph model). The model is built
    -- to reduce the worst-case number of scalar reads; an optimal migration of
    -- statements depends on runtime data.
    --
    -- Blocking scalar writes are reduced by either turning such writes into
    -- asynchronous kernels, as is done with scalar array literals and accumulator
    -- updates, or by transforming host-device writing into device-device copying.
    --
    -- For details on how the graph is constructed and how the vertex cut is found,
    -- see the master thesis "Reducing Synchronous GPU Memory Transfers" by Philip
    -- Børgesen (2022).
    module Futhark.Optimise.ReduceDeviceSyncs.MigrationTable
      ( -- * Analysis
        analyseFunDef,
        analyseConsts,
        hostOnlyFunDefs,

        -- * Types
        MigrationTable,
        MigrationStatus (..),

        -- * Query

        -- | These functions all assume that no parent statement should be migrated.
        -- That is @shouldMoveStm stm mt@ should return @False@ for every statement
        -- @stm@ with a body that a queried 'VName' or 'Stm' is nested within,
        -- otherwise the query result may be invalid.
        shouldMoveStm,
        shouldMove,
        usedOnHost,
        statusOf,
      )
    where

    import Control.Monad
    import Control.Monad.Trans.Class
    import Control.Monad.Trans.Reader qualified as R
    import Control.Monad.Trans.State.Strict ()
    import Control.Monad.Trans.State.Strict hiding (State)
    import Data.Bifunctor (first, second)
    import Data.Foldable
    import Data.IntMap.Strict qualified as IM
    import Data.IntSet qualified as IS
    import Data.List qualified as L
    import Data.Map.Strict qualified as M
    import Data.Maybe (fromMaybe, isJust, isNothing)
    import Data.Sequence qualified as SQ
    import Data.Set (Set, (\\))
    import Data.Set qualified as S
    import Futhark.Error
    import Futhark.IR.GPU
    import Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph
      ( EdgeType (..),
        Edges (..),
        Id,
        IdSet,
        Result (..),
        Routing (..),
        Vertex (..),
      )
    import Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph qualified as MG

    --------------------------------------------------------------------------------
    --                              MIGRATION TABLES                              --
    --------------------------------------------------------------------------------

    -- | Where the value bound by a name should be computed.
    data MigrationStatus
      = -- | The statement that computes the value should be moved to device.
        -- No host usage of the value will be left after the migration.
        MoveToDevice
      | -- | As 'MoveToDevice' but host usage of the value will remain after
        -- migration.
        UsedOnHost
      | -- | The statement that computes the value should remain on host.
        StayOnHost
      deriving (MigrationStatus -> MigrationStatus -> Bool
    (MigrationStatus -> MigrationStatus -> Bool)
    -> (MigrationStatus -> MigrationStatus -> Bool)
    -> Eq MigrationStatus
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: MigrationStatus -> MigrationStatus -> Bool
    == :: MigrationStatus -> MigrationStatus -> Bool
    $c/= :: MigrationStatus -> MigrationStatus -> Bool
    /= :: MigrationStatus -> MigrationStatus -> Bool
    Eq, Eq MigrationStatus
    Eq MigrationStatus =>
    (MigrationStatus -> MigrationStatus -> Ordering)
    -> (MigrationStatus -> MigrationStatus -> Bool)
    -> (MigrationStatus -> MigrationStatus -> Bool)
    -> (MigrationStatus -> MigrationStatus -> Bool)
    -> (MigrationStatus -> MigrationStatus -> Bool)
    -> (MigrationStatus -> MigrationStatus -> MigrationStatus)
    -> (MigrationStatus -> MigrationStatus -> MigrationStatus)
    -> Ord MigrationStatus
    MigrationStatus -> MigrationStatus -> Bool
    MigrationStatus -> MigrationStatus -> Ordering
    MigrationStatus -> MigrationStatus -> MigrationStatus
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: MigrationStatus -> MigrationStatus -> Ordering
    compare :: MigrationStatus -> MigrationStatus -> Ordering
    $c< :: MigrationStatus -> MigrationStatus -> Bool
    < :: MigrationStatus -> MigrationStatus -> Bool
    $c<= :: MigrationStatus -> MigrationStatus -> Bool
    <= :: MigrationStatus -> MigrationStatus -> Bool
    $c> :: MigrationStatus -> MigrationStatus -> Bool
    > :: MigrationStatus -> MigrationStatus -> Bool
    $c>= :: MigrationStatus -> MigrationStatus -> Bool
    >= :: MigrationStatus -> MigrationStatus -> Bool
    $cmax :: MigrationStatus -> MigrationStatus -> MigrationStatus
    max :: MigrationStatus -> MigrationStatus -> MigrationStatus
    $cmin :: MigrationStatus -> MigrationStatus -> MigrationStatus
    min :: MigrationStatus -> MigrationStatus -> MigrationStatus
    Ord, Id -> MigrationStatus -> ShowS
    [MigrationStatus] -> ShowS
    MigrationStatus -> String
    (Id -> MigrationStatus -> ShowS)
    -> (MigrationStatus -> String)
    -> ([MigrationStatus] -> ShowS)
    -> Show MigrationStatus
    forall a.
    (Id -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Id -> MigrationStatus -> ShowS
    showsPrec :: Id -> MigrationStatus -> ShowS
    $cshow :: MigrationStatus -> String
    show :: MigrationStatus -> String
    $cshowList :: [MigrationStatus] -> ShowS
    showList :: [MigrationStatus] -> ShowS
    Show)

    -- | Identifies
    --
    --     (1) which statements should be moved from host to device to reduce the
    --         worst case number of blocking memory transfers.
    --
    --     (2) which migrated variables that still will be used on the host after
    --         all such statements have been moved.
    newtype MigrationTable = MigrationTable (IM.IntMap MigrationStatus)

    instance Semigroup MigrationTable where
      MigrationTable IntMap MigrationStatus
    a <> :: MigrationTable -> MigrationTable -> MigrationTable
    <> MigrationTable IntMap MigrationStatus
    b = IntMap MigrationStatus -> MigrationTable
    MigrationTable (IntMap MigrationStatus
    a IntMap MigrationStatus
    -> IntMap MigrationStatus -> IntMap MigrationStatus
    forall a. IntMap a -> IntMap a -> IntMap a
    `IM.union` IntMap MigrationStatus
    b)

    -- | Where should the value bound by this name be computed?
    statusOf :: VName -> MigrationTable -> MigrationStatus
    statusOf :: VName -> MigrationTable -> MigrationStatus
    statusOf VName
    n (MigrationTable IntMap MigrationStatus
    mt) =
      MigrationStatus -> Maybe MigrationStatus -> MigrationStatus
    forall a. a -> Maybe a -> a
    fromMaybe MigrationStatus
    StayOnHost (Maybe MigrationStatus -> MigrationStatus)
    -> Maybe MigrationStatus -> MigrationStatus
    forall a b. (a -> b) -> a -> b
    $ Id -> IntMap MigrationStatus -> Maybe MigrationStatus
    forall a. Id -> IntMap a -> Maybe a
    IM.lookup (VName -> Id
    baseTag VName
    n) IntMap MigrationStatus
    mt

    -- | Should this whole statement be moved from host to device?
    shouldMoveStm :: Stm GPU -> MigrationTable -> Bool
    shouldMoveStm :: Stm GPU -> MigrationTable -> Bool
    shouldMoveStm (Let (Pat ((PatElem VName
    n LetDec GPU
    _) : [PatElem (LetDec GPU)]
    _)) StmAux (ExpDec GPU)
    _ (BasicOp (Index VName
    _ Slice SubExp
    slice))) MigrationTable
    mt =
      VName -> MigrationTable -> MigrationStatus
    statusOf VName
    n MigrationTable
    mt MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    == MigrationStatus
    MoveToDevice Bool -> Bool -> Bool
    || (SubExp -> Bool) -> Slice SubExp -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any SubExp -> Bool
    movedOperand Slice SubExp
    slice
      where
        movedOperand :: SubExp -> Bool
    movedOperand (Var VName
    op) = VName -> MigrationTable -> MigrationStatus
    statusOf VName
    op MigrationTable
    mt MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    == MigrationStatus
    MoveToDevice
        movedOperand SubExp
    _ = Bool
    False
    shouldMoveStm (Let (Pat ((PatElem VName
    n LetDec GPU
    _) : [PatElem (LetDec GPU)]
    _)) StmAux (ExpDec GPU)
    _ (BasicOp BasicOp
    _)) MigrationTable
    mt =
      VName -> MigrationTable -> MigrationStatus
    statusOf VName
    n MigrationTable
    mt MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    /= MigrationStatus
    StayOnHost
    shouldMoveStm (Let (Pat ((PatElem VName
    n LetDec GPU
    _) : [PatElem (LetDec GPU)]
    _)) StmAux (ExpDec GPU)
    _ Apply {}) MigrationTable
    mt =
      VName -> MigrationTable -> MigrationStatus
    statusOf VName
    n MigrationTable
    mt MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    /= MigrationStatus
    StayOnHost
    shouldMoveStm (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (Match [SubExp]
    cond [Case (Body GPU)]
    _ Body GPU
    _ MatchDec (BranchType GPU)
    _)) MigrationTable
    mt =
      (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ((MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    == MigrationStatus
    MoveToDevice) (MigrationStatus -> Bool)
    -> (VName -> MigrationStatus) -> VName -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName -> MigrationTable -> MigrationStatus
    `statusOf` MigrationTable
    mt)) ([VName] -> Bool) -> [VName] -> Bool
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> [VName]
    subExpVars [SubExp]
    cond
    shouldMoveStm (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (Loop [(FParam GPU, SubExp)]
    _ (ForLoop VName
    _ IntType
    _ (Var VName
    n)) Body GPU
    _)) MigrationTable
    mt =
      VName -> MigrationTable -> MigrationStatus
    statusOf VName
    n MigrationTable
    mt MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    == MigrationStatus
    MoveToDevice
    shouldMoveStm (Let Pat (LetDec GPU)
    _ StmAux (ExpDec GPU)
    _ (Loop [(FParam GPU, SubExp)]
    _ (WhileLoop VName
    n) Body GPU
    _)) MigrationTable
    mt =
      VName -> MigrationTable -> MigrationStatus
    statusOf VName
    n MigrationTable
    mt MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    == MigrationStatus
    MoveToDevice
    -- BasicOp and Apply statements might not bind any variables (shouldn't happen).
    -- If statements might use a constant branch condition.
    -- For loop statements might use a constant number of iterations.
    -- HostOp statements cannot execute on device.
    -- WithAcc statements are never moved in their entirety.
    shouldMoveStm Stm GPU
    _ MigrationTable
    _ = Bool
    False

    -- | Should the value bound by this name be computed on device?
    shouldMove :: VName -> MigrationTable -> Bool
    shouldMove :: VName -> MigrationTable -> Bool
    shouldMove VName
    n MigrationTable
    mt = VName -> MigrationTable -> MigrationStatus
    statusOf VName
    n MigrationTable
    mt MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    /= MigrationStatus
    StayOnHost

    -- | Will the value bound by this name be used on host?
    usedOnHost :: VName -> MigrationTable -> Bool
    usedOnHost :: VName -> MigrationTable -> Bool
    usedOnHost VName
    n MigrationTable
    mt = VName -> MigrationTable -> MigrationStatus
    statusOf VName
    n MigrationTable
    mt MigrationStatus -> MigrationStatus -> Bool
    forall a. Eq a => a -> a -> Bool
    /= MigrationStatus
    MoveToDevice

    --------------------------------------------------------------------------------
    --                         HOST-ONLY FUNCTION ANALYSIS                        --
    --------------------------------------------------------------------------------

    -- | Identifies top-level function definitions that cannot be run on the
    -- device. The application of any such function is host-only.
    type HostOnlyFuns = Set Name

    -- | Returns the names of all top-level functions that cannot be called from the
    -- device. The evaluation of such a function is host-only.
    hostOnlyFunDefs :: [FunDef GPU] -> HostOnlyFuns
    hostOnlyFunDefs :: [FunDef GPU] -> HostOnlyFuns
    hostOnlyFunDefs [FunDef GPU]
    funs =
      let names :: [Name]
    names = (FunDef GPU -> Name) -> [FunDef GPU] -> [Name]
    forall a b. (a -> b) -> [a] -> [b]
    map FunDef GPU -> Name
    forall rep. FunDef rep -> Name
    funDefName [FunDef GPU]
    funs
          call_map :: Map Name (Maybe HostOnlyFuns)
    call_map = [(Name, Maybe HostOnlyFuns)] -> Map Name (Maybe HostOnlyFuns)
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(Name, Maybe HostOnlyFuns)] -> Map Name (Maybe HostOnlyFuns))
    -> [(Name, Maybe HostOnlyFuns)] -> Map Name (Maybe HostOnlyFuns)
    forall a b. (a -> b) -> a -> b
    $ [Name] -> [Maybe HostOnlyFuns] -> [(Name, Maybe HostOnlyFuns)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Name]
    names ((FunDef GPU -> Maybe HostOnlyFuns)
    -> [FunDef GPU] -> [Maybe HostOnlyFuns]
    forall a b. (a -> b) -> [a] -> [b]
    map FunDef GPU -> Maybe HostOnlyFuns
    checkFunDef [FunDef GPU]
    funs)
       in [Name] -> HostOnlyFuns
    forall a. Ord a => [a] -> Set a
    S.fromList [Name]
    names HostOnlyFuns -> HostOnlyFuns -> HostOnlyFuns
    forall a. Ord a => Set a -> Set a -> Set a
    \\ Map Name (Maybe HostOnlyFuns) -> HostOnlyFuns
    forall {a}. Map Name a -> HostOnlyFuns
    keysToSet (Map Name (Maybe HostOnlyFuns) -> Map Name (Maybe HostOnlyFuns)
    removeHostOnly Map Name (Maybe HostOnlyFuns)
    call_map)
      where
        keysToSet :: Map Name a -> HostOnlyFuns
    keysToSet = [Name] -> HostOnlyFuns
    forall a. Eq a => [a] -> Set a
    S.fromAscList ([Name] -> HostOnlyFuns)
    -> (Map Name a -> [Name]) -> Map Name a -> HostOnlyFuns
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map Name a -> [Name]
    forall k a. Map k a -> [k]
    M.keys

        removeHostOnly :: Map Name (Maybe HostOnlyFuns) -> Map Name (Maybe HostOnlyFuns)
    removeHostOnly Map Name (Maybe HostOnlyFuns)
    cm =
          let (Map Name (Maybe HostOnlyFuns)
    host_only, Map Name (Maybe HostOnlyFuns)
    cm') = (Maybe HostOnlyFuns -> Bool)
    -> Map Name (Maybe HostOnlyFuns)
    -> (Map Name (Maybe HostOnlyFuns), Map Name (Maybe HostOnlyFuns))
    forall a k. (a -> Bool) -> Map k a -> (Map k a, Map k a)
    M.partition Maybe HostOnlyFuns -> Bool
    forall {a}. Maybe a -> Bool
    isHostOnly Map Name (Maybe HostOnlyFuns)
    cm
           in if Map Name (Maybe HostOnlyFuns) -> Bool
    forall k a. Map k a -> Bool
    M.null Map Name (Maybe HostOnlyFuns)
    host_only
                then Map Name (Maybe HostOnlyFuns)
    cm'
                else Map Name (Maybe HostOnlyFuns) -> Map Name (Maybe HostOnlyFuns)
    removeHostOnly (Map Name (Maybe HostOnlyFuns) -> Map Name (Maybe HostOnlyFuns))
    -> Map Name (Maybe HostOnlyFuns) -> Map Name (Maybe HostOnlyFuns)
    forall a b. (a -> b) -> a -> b
    $ (Maybe HostOnlyFuns -> Maybe HostOnlyFuns)
    -> Map Name (Maybe HostOnlyFuns) -> Map Name (Maybe HostOnlyFuns)
    forall a b k. (a -> b) -> Map k a -> Map k b
    M.map (HostOnlyFuns -> Maybe HostOnlyFuns -> Maybe HostOnlyFuns
    forall {a}. Ord a => Set a -> Maybe (Set a) -> Maybe (Set a)
    checkCalls (HostOnlyFuns -> Maybe HostOnlyFuns -> Maybe HostOnlyFuns)
    -> HostOnlyFuns -> Maybe HostOnlyFuns -> Maybe HostOnlyFuns
    forall a b. (a -> b) -> a -> b
    $ Map Name (Maybe HostOnlyFuns) -> HostOnlyFuns
    forall {a}. Map Name a -> HostOnlyFuns
    keysToSet Map Name (Maybe HostOnlyFuns)
    host_only) Map Name (Maybe HostOnlyFuns)
    cm'

        isHostOnly :: Maybe a -> Bool
    isHostOnly = Maybe a -> Bool
    forall {a}. Maybe a -> Bool
    isNothing

        -- A function that calls a host-only function is itself host-only.
        checkCalls :: Set a -> Maybe (Set a) -> Maybe (Set a)
    checkCalls Set a
    hostOnlyFuns (Just Set a
    calls)
          | Set a
    hostOnlyFuns Set a -> Set a -> Bool
    forall a. Ord a => Set a -> Set a -> Bool
    `S.disjoint` Set a
    calls =
              Set a -> Maybe (Set a)
    forall a. a -> Maybe a
    Just Set a
    calls
        checkCalls Set a
    _ Maybe (Set a)
    _ =
          Maybe (Set a)
    forall a. Maybe a
    Nothing

    -- | 'checkFunDef' returns 'Nothing' if this function definition uses arrays or
    -- HostOps. Otherwise it returns the names of all applied functions, which may
    -- include user defined functions that could turn out to be host-only.
    checkFunDef :: FunDef GPU -> Maybe (Set Name)
    checkFunDef :: FunDef GPU -> Maybe HostOnlyFuns
    checkFunDef FunDef GPU
    fun = do
      [Param DeclType] -> Maybe ()
    checkFParams ([Param DeclType] -> Maybe ()) -> [Param DeclType] -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ FunDef GPU -> [FParam GPU]
    forall rep. FunDef rep -> [FParam rep]
    funDefParams FunDef GPU
    fun
      [TypeBase ExtShape Uniqueness] -> Maybe ()
    forall {u}. [TypeBase ExtShape u] -> Maybe ()
    checkRetTypes ([TypeBase ExtShape Uniqueness] -> Maybe ())
    -> [TypeBase ExtShape Uniqueness] -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ ((TypeBase ExtShape Uniqueness, RetAls)
     -> TypeBase ExtShape Uniqueness)
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map (TypeBase ExtShape Uniqueness, RetAls)
    -> TypeBase ExtShape Uniqueness
    forall a b. (a, b) -> a
    fst ([(TypeBase ExtShape Uniqueness, RetAls)]
     -> [TypeBase ExtShape Uniqueness])
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> a -> b
    $ FunDef GPU -> [(RetType GPU, RetAls)]
    forall rep. FunDef rep -> [(RetType rep, RetAls)]
    funDefRetType FunDef GPU
    fun
      Body GPU -> Maybe HostOnlyFuns
    checkBody (Body GPU -> Maybe HostOnlyFuns) -> Body GPU -> Maybe HostOnlyFuns
    forall a b. (a -> b) -> a -> b
    $ FunDef GPU -> Body GPU
    forall rep. FunDef rep -> Body rep
    funDefBody FunDef GPU
    fun
      where
        hostOnly :: Maybe a
    hostOnly = Maybe a
    forall a. Maybe a
    Nothing
        ok :: Maybe ()
    ok = () -> Maybe ()
    forall a. a -> Maybe a
    Just ()
        check :: (a -> Bool) -> t a -> Maybe ()
    check a -> Bool
    isArr t a
    as = if (a -> Bool) -> t a -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any a -> Bool
    isArr t a
    as then Maybe ()
    forall a. Maybe a
    hostOnly else Maybe ()
    ok

        checkFParams :: [Param DeclType] -> Maybe ()
    checkFParams = (Param DeclType -> Bool) -> [Param DeclType] -> Maybe ()
    forall {t :: * -> *} {a}.
    Foldable t =>
    (a -> Bool) -> t a -> Maybe ()
    check Param DeclType -> Bool
    forall t. Typed t => t -> Bool
    isArray

        checkLParams :: [(FParam GPU, b)] -> Maybe ()
    checkLParams = ((FParam GPU, b) -> Bool) -> [(FParam GPU, b)] -> Maybe ()
    forall {t :: * -> *} {a}.
    Foldable t =>
    (a -> Bool) -> t a -> Maybe ()
    check (FParam GPU -> Bool
    forall t. Typed t => t -> Bool
    isArray (FParam GPU -> Bool)
    -> ((FParam GPU, b) -> FParam GPU) -> (FParam GPU, b) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (FParam GPU, b) -> FParam GPU
    forall a b. (a, b) -> a
    fst)

        checkRetTypes :: [TypeBase ExtShape u] -> Maybe ()
    checkRetTypes = (TypeBase ExtShape u -> Bool) -> [TypeBase ExtShape u] -> Maybe ()
    forall {t :: * -> *} {a}.
    Foldable t =>
    (a -> Bool) -> t a -> Maybe ()
    check TypeBase ExtShape u -> Bool
    forall shape u. ArrayShape shape => TypeBase shape u -> Bool
    isArrayType

        checkPats :: [PatElem (LetDec GPU)] -> Maybe ()
    checkPats = (PatElem (LetDec GPU) -> Bool)
    -> [PatElem (LetDec GPU)] -> Maybe ()
    forall {t :: * -> *} {a}.
    Foldable t =>
    (a -> Bool) -> t a -> Maybe ()
    check PatElem (LetDec GPU) -> Bool
    forall t. Typed t => t -> Bool
    isArray

        checkBody :: Body GPU -> Maybe HostOnlyFuns
    checkBody = Stms GPU -> Maybe HostOnlyFuns
    checkStms (Stms GPU -> Maybe HostOnlyFuns)
    -> (Body GPU -> Stms GPU) -> Body GPU -> Maybe HostOnlyFuns
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body GPU -> Stms GPU
    forall rep. Body rep -> Stms rep
    bodyStms

        checkStms :: Stms GPU -> Maybe HostOnlyFuns
    checkStms Stms GPU
    stms = Seq HostOnlyFuns -> HostOnlyFuns
    forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
    S.unions (Seq HostOnlyFuns -> HostOnlyFuns)
    -> Maybe (Seq HostOnlyFuns) -> Maybe HostOnlyFuns
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Stm GPU -> Maybe HostOnlyFuns)
    -> Stms GPU -> Maybe (Seq HostOnlyFuns)
    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) -> Seq a -> m (Seq b)
    mapM Stm GPU -> Maybe HostOnlyFuns
    checkStm Stms GPU
    stms

        checkStm :: Stm GPU -> Maybe HostOnlyFuns
    checkStm (Let (Pat [PatElem (LetDec GPU)]
    pats) StmAux (ExpDec GPU)
    _ Exp GPU
    e) = [PatElem (LetDec GPU)] -> Maybe ()
    checkPats [PatElem (LetDec GPU)]
    pats Maybe () -> Maybe HostOnlyFuns -> Maybe HostOnlyFuns
    forall a b. Maybe a -> Maybe b -> Maybe b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Exp GPU -> Maybe HostOnlyFuns
    checkExp Exp GPU
    e

        -- Any expression that produces an array is caught by checkPats
        checkExp :: Exp GPU -> Maybe HostOnlyFuns
    checkExp (BasicOp (Index VName
    _ Slice SubExp
    _)) = Maybe HostOnlyFuns
    forall a. Maybe a
    hostOnly
        checkExp (WithAcc [WithAccInput GPU]
    _ Lambda GPU
    _) = Maybe HostOnlyFuns
    forall a. Maybe a
    hostOnly
        checkExp (Op Op GPU
    _) = Maybe HostOnlyFuns
    forall a. Maybe a
    hostOnly
        checkExp (Apply Name
    fn [(SubExp, Diet)]
    _ [(RetType GPU, RetAls)]
    _ Safety
    _) = HostOnlyFuns -> Maybe HostOnlyFuns
    forall a. a -> Maybe a
    Just (Name -> HostOnlyFuns
    forall a. a -> Set a
    S.singleton Name
    fn)
        checkExp (Match [SubExp]
    _ [Case (Body GPU)]
    cases Body GPU
    defbody MatchDec (BranchType GPU)
    _) =
          [HostOnlyFuns] -> HostOnlyFuns
    forall a. Monoid a => [a] -> a
    mconcat ([HostOnlyFuns] -> HostOnlyFuns)
    -> Maybe [HostOnlyFuns] -> Maybe HostOnlyFuns
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Body GPU -> Maybe HostOnlyFuns)
    -> [Body GPU] -> Maybe [HostOnlyFuns]
    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 Body GPU -> Maybe HostOnlyFuns
    checkBody (Body GPU
    defbody Body GPU -> [Body GPU] -> [Body GPU]
    forall a. a -> [a] -> [a]
    : (Case (Body GPU) -> Body GPU) -> [Case (Body GPU)] -> [Body GPU]
    forall a b. (a -> b) -> [a] -> [b]
    map Case (Body GPU) -> Body GPU
    forall body. Case body -> body
    caseBody [Case (Body GPU)]
    cases)
        checkExp (Loop [(FParam GPU, SubExp)]
    params LoopForm
    _ Body GPU
    body) = do
          [(FParam GPU, SubExp)] -> Maybe ()
    forall {b}. [(FParam GPU, b)] -> Maybe ()
    checkLParams [(FParam GPU, SubExp)]
    params
          Body GPU -> Maybe HostOnlyFuns
    checkBody Body GPU
    body
        checkExp BasicOp {} = HostOnlyFuns -> Maybe HostOnlyFuns
    forall a. a -> Maybe a
    Just HostOnlyFuns
    forall a. Set a
    S.empty

    --------------------------------------------------------------------------------
    --                             MIGRATION ANALYSIS                             --
    --------------------------------------------------------------------------------

    -- | HostUsage identifies scalar variables that are used on host.
    type HostUsage = [Id]

    nameToId :: VName -> Id
    nameToId :: VName -> Id
    nameToId = VName -> Id
    baseTag

    -- | Analyses top-level constants.
    analyseConsts :: HostOnlyFuns -> [FunDef GPU] -> Stms GPU -> MigrationTable
    analyseConsts :: HostOnlyFuns -> [FunDef GPU] -> Stms GPU -> MigrationTable
    analyseConsts HostOnlyFuns
    hof [FunDef GPU]
    funs Stms GPU
    consts =
      let usage :: [Id]
    usage = ([Id] -> VName -> NameInfo GPU -> [Id])
    -> [Id] -> Map VName (NameInfo GPU) -> [Id]
    forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
    M.foldlWithKey (Names -> [Id] -> VName -> NameInfo GPU -> [Id]
    forall {t}. Typed t => Names -> [Id] -> VName -> t -> [Id]
    f (Names -> [Id] -> VName -> NameInfo GPU -> [Id])
    -> Names -> [Id] -> VName -> NameInfo GPU -> [Id]
    forall a b. (a -> b) -> a -> b
    $ [FunDef GPU] -> Names
    forall a. FreeIn a => a -> Names
    freeIn [FunDef GPU]
    funs) [] (Stms GPU -> Map VName (NameInfo GPU)
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf Stms GPU
    consts)
       in HostOnlyFuns -> [Id] -> Stms GPU -> MigrationTable
    analyseStms HostOnlyFuns
    hof [Id]
    usage Stms GPU
    consts
      where
        f :: Names -> [Id] -> VName -> t -> [Id]
    f Names
    free [Id]
    usage VName
    n t
    t
          | t -> Bool
    forall t. Typed t => t -> Bool
    isScalar t
    t,
            VName
    n VName -> Names -> Bool
    `nameIn` Names
    free =
              VName -> Id
    nameToId VName
    n Id -> [Id] -> [Id]
    forall a. a -> [a] -> [a]
    : [Id]
    usage
          | Bool
    otherwise =
              [Id]
    usage

    -- | Analyses a top-level function definition.
    analyseFunDef :: HostOnlyFuns -> FunDef GPU -> MigrationTable
    analyseFunDef :: HostOnlyFuns -> FunDef GPU -> MigrationTable
    analyseFunDef HostOnlyFuns
    hof FunDef GPU
    fd =
      let body :: Body GPU
    body = FunDef GPU -> Body GPU
    forall rep. FunDef rep -> Body rep
    funDefBody FunDef GPU
    fd
          usage :: [Id]
    usage = ([Id] -> (SubExpRes, TypeBase ExtShape Uniqueness) -> [Id])
    -> [Id] -> [(SubExpRes, TypeBase ExtShape Uniqueness)] -> [Id]
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' [Id] -> (SubExpRes, TypeBase ExtShape Uniqueness) -> [Id]
    forall {shape} {u}. [Id] -> (SubExpRes, TypeBase shape u) -> [Id]
    f [] ([(SubExpRes, TypeBase ExtShape Uniqueness)] -> [Id])
    -> [(SubExpRes, TypeBase ExtShape Uniqueness)] -> [Id]
    forall a b. (a -> b) -> a -> b
    $ [SubExpRes]
    -> [TypeBase ExtShape Uniqueness]
    -> [(SubExpRes, TypeBase ExtShape Uniqueness)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Body GPU -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult Body GPU
    body) (((TypeBase ExtShape Uniqueness, RetAls)
     -> TypeBase ExtShape Uniqueness)
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> [a] -> [b]
    map (TypeBase ExtShape Uniqueness, RetAls)
    -> TypeBase ExtShape Uniqueness
    forall a b. (a, b) -> a
    fst ([(TypeBase ExtShape Uniqueness, RetAls)]
     -> [TypeBase ExtShape Uniqueness])
    -> [(TypeBase ExtShape Uniqueness, RetAls)]
    -> [TypeBase ExtShape Uniqueness]
    forall a b. (a -> b) -> a -> b
    $ FunDef GPU -> [(RetType GPU, RetAls)]
    forall rep. FunDef rep -> [(RetType rep, RetAls)]
    funDefRetType FunDef GPU
    fd)
          stms :: Stms GPU
    stms = Body GPU -> Stms GPU
    forall rep. Body rep -> Stms rep
    bodyStms Body GPU
    body
       in HostOnlyFuns -> [Id] -> Stms GPU -> MigrationTable
    analyseStms HostOnlyFuns
    hof [Id]
    usage Stms GPU
    stms
      where
        f :: [Id] -> (SubExpRes, TypeBase shape u) -> [Id]
    f [Id]
    usage (SubExpRes Certs
    _ (Var VName
    n), TypeBase shape u
    t) | TypeBase shape u -> Bool
    forall shape u. TypeBase shape u -> Bool
    isScalarType TypeBase shape u
    t = VName -> Id
    nameToId VName
    n Id -> [Id] -> [Id]
    forall a. a -> [a] -> [a]
    : [Id]
    usage
        f [Id]
    usage (SubExpRes, TypeBase shape u)
    _ = [Id]
    usage

    -- | Analyses statements. The 'HostUsage' list identifies which bound scalar
    -- variables that subsequently may be used on host. All free variables such as
    -- constants and function parameters are assumed to reside on host.
    analyseStms :: HostOnlyFuns -> HostUsage -> Stms GPU -> MigrationTable
    analyseStms :: HostOnlyFuns -> [Id] -> Stms GPU -> MigrationTable
    analyseStms HostOnlyFuns
    hof [Id]
    usage Stms GPU
    stms =
      let (Graph
    g, Sources
    srcs, [Id]
    _) = HostOnlyFuns -> [Id] -> Stms GPU -> (Graph, Sources, [Id])
    buildGraph HostOnlyFuns
    hof [Id]
    usage Stms GPU
    stms
          ([Id]
    routed, [Id]
    unrouted) = Sources
    srcs
          ([Id]
    _, Graph
    g') = [Id] -> Graph -> ([Id], Graph)
    forall m. [Id] -> Graph m -> ([Id], Graph m)
    MG.routeMany [Id]
    unrouted Graph
    g -- hereby routed
          f :: ((Operands, Operands, Operands), Visited ())
    -> Id -> ((Operands, Operands, Operands), Visited ())
    f ((Operands, Operands, Operands), Visited ())
    st' = Graph
    -> ((Operands, Operands, Operands)
        -> EdgeType -> Vertex Meta -> (Operands, Operands, Operands))
    -> ((Operands, Operands, Operands), Visited ())
    -> EdgeType
    -> Id
    -> ((Operands, Operands, Operands), Visited ())
    forall m a.
    Graph m
    -> (a -> EdgeType -> Vertex m -> a)
    -> (a, Visited ())
    -> EdgeType
    -> Id
    -> (a, Visited ())
    MG.fold Graph
    g' (Operands, Operands, Operands)
    -> EdgeType -> Vertex Meta -> (Operands, Operands, Operands)
    forall {m}.
    (Operands, Operands, Operands)
    -> EdgeType -> Vertex m -> (Operands, Operands, Operands)
    visit ((Operands, Operands, Operands), Visited ())
    st' EdgeType
    Normal
          st :: ((Operands, Operands, Operands), Visited ())
    st = (((Operands, Operands, Operands), Visited ())
     -> Id -> ((Operands, Operands, Operands), Visited ()))
    -> ((Operands, Operands, Operands), Visited ())
    -> [Id]
    -> ((Operands, Operands, Operands), Visited ())
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' ((Operands, Operands, Operands), Visited ())
    -> Id -> ((Operands, Operands, Operands), Visited ())
    f ((Operands, Operands, Operands)
    initial, Visited ()
    forall a. Visited a
    MG.none) [Id]
    unrouted
          (Operands
    vr, Operands
    vn, Operands
    tn) = ((Operands, Operands, Operands), Visited ())
    -> (Operands, Operands, Operands)
    forall a b. (a, b) -> a
    fst (((Operands, Operands, Operands), Visited ())
     -> (Operands, Operands, Operands))
    -> ((Operands, Operands, Operands), Visited ())
    -> (Operands, Operands, Operands)
    forall a b. (a -> b) -> a -> b
    $ (((Operands, Operands, Operands), Visited ())
     -> Id -> ((Operands, Operands, Operands), Visited ()))
    -> ((Operands, Operands, Operands), Visited ())
    -> [Id]
    -> ((Operands, Operands, Operands), Visited ())
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' ((Operands, Operands, Operands), Visited ())
    -> Id -> ((Operands, Operands, Operands), Visited ())
    f ((Operands, Operands, Operands), Visited ())
    st [Id]
    routed
       in -- TODO: Delay reads into (deeper) branches

          IntMap MigrationStatus -> MigrationTable
    MigrationTable (IntMap MigrationStatus -> MigrationTable)
    -> IntMap MigrationStatus -> MigrationTable
    forall a b. (a -> b) -> a -> b
    $
            [IntMap MigrationStatus] -> IntMap MigrationStatus
    forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
    IM.unions
              [ (Id -> MigrationStatus) -> Operands -> IntMap MigrationStatus
    forall a. (Id -> a) -> Operands -> IntMap a
    IM.fromSet (MigrationStatus -> Id -> MigrationStatus
    forall a b. a -> b -> a
    const MigrationStatus
    MoveToDevice) Operands
    vr,
                (Id -> MigrationStatus) -> Operands -> IntMap MigrationStatus
    forall a. (Id -> a) -> Operands -> IntMap a
    IM.fromSet (MigrationStatus -> Id -> MigrationStatus
    forall a b. a -> b -> a
    const MigrationStatus
    MoveToDevice) Operands
    vn,
                -- Read by host if not reached by a reversed edge
                (Id -> MigrationStatus) -> Operands -> IntMap MigrationStatus
    forall a. (Id -> a) -> Operands -> IntMap a
    IM.fromSet (MigrationStatus -> Id -> MigrationStatus
    forall a b. a -> b -> a
    const MigrationStatus
    UsedOnHost) Operands
    tn
              ]
      where
        -- 1) Visited by reversed edge.
        -- 2) Visited by normal edge, no route.
        -- 3) Visited by normal edge, had route; will potentially be read by host.
        initial :: (Operands, Operands, Operands)
    initial = (Operands
    IS.empty, Operands
    IS.empty, Operands
    IS.empty)

        visit :: (Operands, Operands, Operands)
    -> EdgeType -> Vertex m -> (Operands, Operands, Operands)
    visit (Operands
    vr, Operands
    vn, Operands
    tn) EdgeType
    Reversed Vertex m
    v =
          let vr' :: Operands
    vr' = Id -> Operands -> Operands
    IS.insert (Vertex m -> Id
    forall m. Vertex m -> Id
    vertexId Vertex m
    v) Operands
    vr
           in (Operands
    vr', Operands
    vn, Operands
    tn)
        visit (Operands
    vr, Operands
    vn, Operands
    tn) EdgeType
    Normal v :: Vertex m
    v@Vertex {vertexRouting :: forall m. Vertex m -> Routing
    vertexRouting = Routing
    NoRoute} =
          let vn' :: Operands
    vn' = Id -> Operands -> Operands
    IS.insert (Vertex m -> Id
    forall m. Vertex m -> Id
    vertexId Vertex m
    v) Operands
    vn
           in (Operands
    vr, Operands
    vn', Operands
    tn)
        visit (Operands
    vr, Operands
    vn, Operands
    tn) EdgeType
    Normal Vertex m
    v =
          let tn' :: Operands
    tn' = Id -> Operands -> Operands
    IS.insert (Vertex m -> Id
    forall m. Vertex m -> Id
    vertexId Vertex m
    v) Operands
    tn
           in (Operands
    vr, Operands
    vn, Operands
    tn')

    --------------------------------------------------------------------------------
    --                                TYPE HELPERS                                --
    --------------------------------------------------------------------------------

    isScalar :: (Typed t) => t -> Bool
    isScalar :: forall t. Typed t => t -> Bool
    isScalar = Type -> Bool
    forall shape u. TypeBase shape u -> Bool
    isScalarType (Type -> Bool) -> (t -> Type) -> t -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . t -> Type
    forall t. Typed t => t -> Type
    typeOf

    isScalarType :: TypeBase shape u -> Bool
    isScalarType :: forall shape u. TypeBase shape u -> Bool
    isScalarType (Prim PrimType
    Unit) = Bool
    False
    isScalarType (Prim PrimType
    _) = Bool
    True
    isScalarType TypeBase shape u
    _ = Bool
    False

    isArray :: (Typed t) => t -> Bool
    isArray :: forall t. Typed t => t -> Bool
    isArray = Type -> Bool
    forall shape u. ArrayShape shape => TypeBase shape u -> Bool
    isArrayType (Type -> Bool) -> (t -> Type) -> t -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . t -> Type
    forall t. Typed t => t -> Type
    typeOf

    isArrayType :: (ArrayShape shape) => TypeBase shape u -> Bool
    isArrayType :: forall shape u. ArrayShape shape => TypeBase shape u -> Bool
    isArrayType = (Id
    0 <) (Id -> Bool)
    -> (TypeBase shape u -> Id) -> TypeBase shape u -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TypeBase shape u -> Id
    forall shape u. ArrayShape shape => TypeBase shape u -> Id
    arrayRank

    --------------------------------------------------------------------------------
    --                               GRAPH BUILDING                               --
    --------------------------------------------------------------------------------

    buildGraph :: HostOnlyFuns -> HostUsage -> Stms GPU -> (Graph, Sources, Sinks)
    buildGraph :: HostOnlyFuns -> [Id] -> Stms GPU -> (Graph, Sources, [Id])
    buildGraph HostOnlyFuns
    hof [Id]
    usage Stms GPU
    stms =
      let (Graph
    g, Sources
    srcs, [Id]
    sinks) = HostOnlyFuns -> Grapher () -> (Graph, Sources, [Id])
    forall a. HostOnlyFuns -> Grapher a -> (Graph, Sources, [Id])
    execGrapher HostOnlyFuns
    hof (Stms GPU -> Grapher ()
    graphStms Stms GPU
    stms)
          g' :: Graph
    g' = (Graph -> Id -> Graph) -> Graph -> [Id] -> Graph
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' ((Id -> Graph -> Graph) -> Graph -> Id -> Graph
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Id -> Graph -> Graph
    forall m. Id -> Graph m -> Graph m
    MG.connectToSink) Graph
    g [Id]
    usage
       in (Graph
    g', Sources
    srcs, [Id]
    sinks)

    -- | Graph a body.
    graphBody :: Body GPU -> Grapher ()
    graphBody :: Body GPU -> Grapher ()
    graphBody Body GPU
    body = do
      let res_ops :: Operands
    res_ops = Names -> Operands
    namesIntSet (Names -> Operands) -> Names -> Operands
    forall a b. (a -> b) -> a -> b
    $ [SubExpRes] -> Names
    forall a. FreeIn a => a -> Names
    freeIn (Body GPU -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult Body GPU
    body)
      BodyStats
    body_stats <-
        Grapher () -> Grapher BodyStats
    forall a. Grapher a -> Grapher BodyStats
    captureBodyStats (Grapher () -> Grapher BodyStats)
    -> Grapher () -> Grapher BodyStats
    forall a b. (a -> b) -> a -> b
    $
          Grapher () -> Grapher ()
    forall a. Grapher a -> Grapher a
    incBodyDepthFor (Stms GPU -> Grapher ()
    graphStms (Body GPU -> Stms GPU
    forall rep. Body rep -> Stms rep
    bodyStms Body GPU
    body) Grapher () -> Grapher () -> Grapher ()
    forall a b.
    StateT State (Reader Env) a
    -> StateT State (Reader Env) b -> StateT State (Reader Env) b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Operands -> Grapher ()
    tellOperands Operands
    res_ops)

      Id
    body_depth <- (Id
    1 +) (Id -> Id)
    -> StateT State (Reader Env) Id -> StateT State (Reader Env) Id
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> StateT State (Reader Env) Id
    getBodyDepth
      let host_only :: Bool
    host_only = Id -> Operands -> Bool
    IS.member Id
    body_depth (BodyStats -> Operands
    bodyHostOnlyParents BodyStats
    body_stats)
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st ->
        let stats :: BodyStats
    stats = State -> BodyStats
    stateStats State
    st
            hops' :: Operands
    hops' = Id -> Operands -> Operands
    IS.delete Id
    body_depth (BodyStats -> Operands
    bodyHostOnlyParents BodyStats
    stats)
            -- If body contains a variable that is required on host the parent
            -- statement that contains this body cannot be migrated as a whole.
            stats' :: BodyStats
    stats' = if Bool
    host_only then BodyStats
    stats {bodyHostOnly = True} else BodyStats
    stats
         in State
    st {stateStats = stats' {bodyHostOnlyParents = hops'}}

    -- | Graph multiple statements.
    graphStms :: Stms GPU -> Grapher ()
    graphStms :: Stms GPU -> Grapher ()
    graphStms = (Stm GPU -> Grapher ()) -> Stms GPU -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ Stm GPU -> Grapher ()
    graphStm

    -- | Graph a single statement.
    graphStm :: Stm GPU -> Grapher ()
    graphStm :: Stm GPU -> Grapher ()
    graphStm Stm GPU
    stm = do
      let bs :: [Binding]
    bs = Stm GPU -> [Binding]
    boundBy Stm GPU
    stm
      let e :: Exp GPU
    e = Stm GPU -> Exp GPU
    forall rep. Stm rep -> Exp rep
    stmExp Stm GPU
    stm
      -- IMPORTANT! It is generally assumed that all scalars within types and
      -- shapes are present on host. Any expression of a type wherein one of its
      -- scalar operands appears must therefore ensure that that scalar operand is
      -- marked as a size variable (see the 'hostSize' function).
      case Exp GPU
    e of
        BasicOp (SubExp SubExp
    se) -> do
          [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e
          [Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs Binding -> SubExp -> Grapher ()
    `reusesSubExp` SubExp
    se
        BasicOp (Opaque OpaqueOp
    _ SubExp
    se) -> do
          [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e
          [Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs Binding -> SubExp -> Grapher ()
    `reusesSubExp` SubExp
    se
        BasicOp (ArrayLit [SubExp]
    arr Type
    t)
          | Type -> Bool
    forall t. Typed t => t -> Bool
    isScalar Type
    t,
            (SubExp -> Bool) -> [SubExp] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (Maybe VName -> Bool
    forall {a}. Maybe a -> Bool
    isJust (Maybe VName -> Bool) -> (SubExp -> Maybe VName) -> SubExp -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExp -> Maybe VName
    subExpVar) [SubExp]
    arr ->
              -- Migrating an array literal with free variables saves a write for
              -- every scalar it contains. Under some backends the compiler
              -- generates asynchronous writes for scalar constants but otherwise
              -- each write will be synchronous. If all scalars are constants then
              -- the compiler generates more efficient code that copies static
              -- device memory.
              Binding -> Grapher ()
    graphAutoMove ([Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs)
        BasicOp UnOp {} -> [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e
        BasicOp BinOp {} -> [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e
        BasicOp CmpOp {} -> [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e
        BasicOp ConvOp {} -> [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e
        BasicOp Assert {} ->
          -- == OpenCL =============================================================
          --
          -- The next read after the execution of a kernel containing an assertion
          -- will be made asynchronous, followed by an asynchronous read to check
          -- if any assertion failed. The runtime will then block for all enqueued
          -- operations to finish.
          --
          -- Since an assertion only binds a certificate of unit type, an assertion
          -- cannot increase the number of (read) synchronizations that occur. In
          -- this regard it is free to migrate. The synchronization that does occur
          -- is however (presumably) more expensive as the pipeline of GPU work will
          -- be flushed.
          --
          -- Since this cost is difficult to quantify and amortize over assertion
          -- migration candidates (cost depends on ordering of kernels and reads) we
          -- assume it is insignificant. This will likely hold for a system where
          -- multiple threads or processes schedules GPU work, as system-wide
          -- throughput only will decrease if the GPU utilization decreases as a
          -- result.
          --
          -- == CUDA ===============================================================
          --
          -- Under the CUDA backend every read is synchronous and is followed by
          -- a full synchronization that blocks for all enqueued operations to
          -- finish. If any enqueued kernel contained an assertion, another
          -- synchronous read is then made to check if an assertion failed.
          --
          -- Migrating an assertion to save a read may thus introduce new reads, and
          -- the total number of reads can hence either decrease, remain the same,
          -- or even increase, subject to the ordering of reads and kernels that
          -- perform assertions.
          --
          -- Since it is possible to implement the same failure checking scheme as
          -- OpenCL using asynchronous reads (and doing so would be a good idea!)
          -- we consider this to be acceptable.
          --
          -- TODO: Implement the OpenCL failure checking scheme under CUDA. This
          --       should reduce the number of synchronizations per read to one.
          [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e
        BasicOp (Index VName
    _ Slice SubExp
    slice)
          | Slice SubExp -> Bool
    forall {d}. Slice d -> Bool
    isFixed Slice SubExp
    slice ->
              Binding -> Grapher ()
    graphRead ([Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs)
        BasicOp {}
          | [(Id
    _, Type
    t)] <- [Binding]
    bs,
            [SubExp]
    dims <- Type -> [SubExp]
    forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
    arrayDims Type
    t,
            [SubExp]
    dims [SubExp] -> [SubExp] -> Bool
    forall a. Eq a => a -> a -> Bool
    /= [], -- i.e. produces an array
            (SubExp -> Bool) -> [SubExp] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (SubExp -> SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1) [SubExp]
    dims ->
              -- An expression that produces an array that only contains a single
              -- primitive value is as efficient to compute and copy as a scalar,
              -- and introduces no size variables.
              --
              -- This is an exception to the inefficiency rules that comes next.
              [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e
        -- Expressions with a cost sublinear to the size of their result arrays are
        -- risky to migrate as we cannot guarantee that their results are not
        -- returned from a GPUBody, which always copies its return values. Since
        -- this would make the effective asymptotic cost of such statements linear
        -- we block them from being migrated on their own.
        --
        -- The parent statement of an enclosing body may still be migrated as a
        -- whole given that each of its returned arrays either
        --   1) is backed by memory used by a migratable statement within its body.
        --   2) contains just a single element.
        -- An array matching either criterion is denoted "copyable memory" because
        -- the asymptotic cost of copying it is less than or equal to the statement
        -- that produced it. This makes the parent of statements with sublinear cost
        -- safe to migrate.
        BasicOp (Index VName
    arr Slice SubExp
    s) -> do
          [SubExp] -> Exp GPU -> Grapher ()
    forall {t :: * -> *}.
    Foldable t =>
    t SubExp -> Exp GPU -> Grapher ()
    graphInefficientReturn (Slice SubExp -> [SubExp]
    forall d. Slice d -> [d]
    sliceDims Slice SubExp
    s) Exp GPU
    e
          [Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs Binding -> VName -> Grapher ()
    `reuses` VName
    arr
        BasicOp (Update Safety
    _ VName
    arr Slice SubExp
    slice SubExp
    _)
          | Slice SubExp -> Bool
    forall {d}. Slice d -> Bool
    isFixed Slice SubExp
    slice -> do
              [SubExp] -> Exp GPU -> Grapher ()
    forall {t :: * -> *}.
    Foldable t =>
    t SubExp -> Exp GPU -> Grapher ()
    graphInefficientReturn [] Exp GPU
    e
              [Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs Binding -> VName -> Grapher ()
    `reuses` VName
    arr
        BasicOp (FlatIndex VName
    arr FlatSlice SubExp
    s) -> do
          -- Migrating a FlatIndex leads to a memory allocation error.
          --
          -- TODO: Fix FlatIndex memory allocation error.
          --
          -- Can be replaced with 'graphHostOnly e' to disable migration.
          -- A fix can be verified by enabling tests/migration/reuse2_flatindex.fut
          [SubExp] -> Exp GPU -> Grapher ()
    forall {t :: * -> *}.
    Foldable t =>
    t SubExp -> Exp GPU -> Grapher ()
    graphInefficientReturn (FlatSlice SubExp -> [SubExp]
    forall d. FlatSlice d -> [d]
    flatSliceDims FlatSlice SubExp
    s) Exp GPU
    e
          [Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs Binding -> VName -> Grapher ()
    `reuses` VName
    arr
        BasicOp (FlatUpdate VName
    arr FlatSlice SubExp
    _ VName
    _) -> do
          [SubExp] -> Exp GPU -> Grapher ()
    forall {t :: * -> *}.
    Foldable t =>
    t SubExp -> Exp GPU -> Grapher ()
    graphInefficientReturn [] Exp GPU
    e
          [Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs Binding -> VName -> Grapher ()
    `reuses` VName
    arr
        BasicOp (Scratch PrimType
    _ [SubExp]
    s) ->
          -- Migrating a Scratch leads to a memory allocation error.
          --
          -- TODO: Fix Scratch memory allocation error.
          --
          -- Can be replaced with 'graphHostOnly e' to disable migration.
          -- A fix can be verified by enabling tests/migration/reuse4_scratch.fut
          [SubExp] -> Exp GPU -> Grapher ()
    forall {t :: * -> *}.
    Foldable t =>
    t SubExp -> Exp GPU -> Grapher ()
    graphInefficientReturn [SubExp]
    s Exp GPU
    e
        BasicOp (Reshape VName
    arr NewShape SubExp
    s) -> do
          [SubExp] -> Exp GPU -> Grapher ()
    forall {t :: * -> *}.
    Foldable t =>
    t SubExp -> Exp GPU -> Grapher ()
    graphInefficientReturn (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims (ShapeBase SubExp -> [SubExp]) -> ShapeBase SubExp -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ NewShape SubExp -> ShapeBase SubExp
    forall d. NewShape d -> ShapeBase d
    newShape NewShape SubExp
    s) Exp GPU
    e
          [Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs Binding -> VName -> Grapher ()
    `reuses` VName
    arr
        BasicOp (Rearrange VName
    arr [Id]
    _) -> do
          [SubExp] -> Exp GPU -> Grapher ()
    forall {t :: * -> *}.
    Foldable t =>
    t SubExp -> Exp GPU -> Grapher ()
    graphInefficientReturn [] Exp GPU
    e
          [Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs Binding -> VName -> Grapher ()
    `reuses` VName
    arr
        -- Expressions with a cost linear to the size of their result arrays are
        -- inefficient to migrate into GPUBody kernels as such kernels are single-
        -- threaded. For sufficiently large arrays the cost may exceed what is saved
        -- by avoiding reads. We therefore also block these from being migrated,
        -- as well as their parents.
        BasicOp ArrayLit {} ->
          -- An array literal purely of primitive constants can be hoisted out to be
          -- a top-level constant, unless it is to be returned or consumed.
          -- Otherwise its runtime implementation will copy a precomputed static
          -- array and thus behave like a 'Copy'.
          -- Whether the rows are primitive constants or arrays, without any scalar
          -- variable operands such ArrayLit cannot directly prevent a scalar read.
          Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
        BasicOp ArrayVal {} ->
          -- As above.
          Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
        BasicOp Update {} ->
          Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
        BasicOp Concat {} ->
          -- Is unlikely to prevent a scalar read as the only SubExp operand in
          -- practice is a computation of host-only size variables.
          Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
        BasicOp Manifest {} ->
          -- Takes no scalar operands so cannot directly prevent a scalar read.
          -- It is introduced as part of the BlkRegTiling kernel optimization and
          -- is thus unlikely to prevent the migration of a parent which was not
          -- already blocked by some host-only operation.
          Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
        BasicOp Iota {} -> Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
        BasicOp Replicate {} -> Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
        -- END
        BasicOp UpdateAcc {} ->
          Binding -> Exp GPU -> Grapher ()
    graphUpdateAcc ([Binding] -> Binding
    forall {a}. [a] -> a
    one [Binding]
    bs) Exp GPU
    e
        Apply Name
    fn [(SubExp, Diet)]
    _ [(RetType GPU, RetAls)]
    _ Safety
    _ ->
          Name -> [Binding] -> Exp GPU -> Grapher ()
    graphApply Name
    fn [Binding]
    bs Exp GPU
    e
        Match [SubExp]
    ses [Case (Body GPU)]
    cases Body GPU
    defbody MatchDec (BranchType GPU)
    _ ->
          [Binding]
    -> [SubExp] -> [Case (Body GPU)] -> Body GPU -> Grapher ()
    graphMatch [Binding]
    bs [SubExp]
    ses [Case (Body GPU)]
    cases Body GPU
    defbody
        Loop [(FParam GPU, SubExp)]
    params LoopForm
    lform Body GPU
    body ->
          [Binding]
    -> [(FParam GPU, SubExp)] -> LoopForm -> Body GPU -> Grapher ()
    graphLoop [Binding]
    bs [(FParam GPU, SubExp)]
    params LoopForm
    lform Body GPU
    body
        WithAcc [WithAccInput GPU]
    inputs Lambda GPU
    f ->
          [Binding] -> [WithAccInput GPU] -> Lambda GPU -> Grapher ()
    graphWithAcc [Binding]
    bs [WithAccInput GPU]
    inputs Lambda GPU
    f
        Op GPUBody {} ->
          -- A GPUBody can be migrated into a parent GPUBody by replacing it with
          -- its body statements and binding its return values inside 'ArrayLit's.
          Grapher ()
    tellGPUBody
        Op Op GPU
    _ ->
          Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
      where
        one :: [a] -> a
    one [a
    x] = a
    x
        one [a]
    _ = String -> a
    forall a. String -> a
    compilerBugS String
    "Type error: unexpected number of pattern elements."

        isFixed :: Slice d -> Bool
    isFixed = Maybe [d] -> Bool
    forall {a}. Maybe a -> Bool
    isJust (Maybe [d] -> Bool) -> (Slice d -> Maybe [d]) -> Slice d -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Slice d -> Maybe [d]
    forall d. Slice d -> Maybe [d]
    sliceIndices

        -- new_dims may introduce new size variables which must be present on host
        -- when this expression is evaluated.
        graphInefficientReturn :: t SubExp -> Exp GPU -> Grapher ()
    graphInefficientReturn t SubExp
    new_dims Exp GPU
    e = do
          (SubExp -> Grapher ()) -> t SubExp -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SubExp -> Grapher ()
    hostSize t SubExp
    new_dims
          Exp GPU -> Grapher Operands
    graphedScalarOperands Exp GPU
    e Grapher Operands -> (Operands -> Grapher ()) -> Grapher ()
    forall a b.
    StateT State (Reader Env) a
    -> (a -> StateT State (Reader Env) b)
    -> StateT State (Reader Env) b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Edges -> Operands -> Grapher ()
    addEdges Edges
    ToSink

        hostSize :: SubExp -> Grapher ()
    hostSize (Var VName
    n) = VName -> Grapher ()
    hostSizeVar VName
    n
        hostSize SubExp
    _ = () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

        hostSizeVar :: VName -> Grapher ()
    hostSizeVar = Id -> Grapher ()
    requiredOnHost (Id -> Grapher ()) -> (VName -> Id) -> VName -> Grapher ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> Id
    nameToId

    -- | Bindings for all pattern elements bound by a statement.
    boundBy :: Stm GPU -> [Binding]
    boundBy :: Stm GPU -> [Binding]
    boundBy = (PatElem Type -> Binding) -> [PatElem Type] -> [Binding]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(PatElem VName
    n Type
    t) -> (VName -> Id
    nameToId VName
    n, Type
    t)) ([PatElem Type] -> [Binding])
    -> (Stm GPU -> [PatElem Type]) -> Stm GPU -> [Binding]
    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

    -- | Graph a statement which in itself neither reads scalars from device memory
    -- nor forces such scalars to be available on host. Such statement can be moved
    -- to device to eliminate the host usage of its operands which transitively may
    -- depend on a scalar device read.
    graphSimple :: [Binding] -> Exp GPU -> Grapher ()
    graphSimple :: [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e = do
      -- Only add vertices to the graph if they have a transitive dependency to
      -- an array read. Transitive dependencies through variables connected to
      -- sinks do not count.
      Operands
    ops <- Exp GPU -> Grapher Operands
    graphedScalarOperands Exp GPU
    e
      let edges :: Edges
    edges = [Id] -> Edges
    MG.declareEdges ((Binding -> Id) -> [Binding] -> [Id]
    forall a b. (a -> b) -> [a] -> [b]
    map Binding -> Id
    forall a b. (a, b) -> a
    fst [Binding]
    bs)
      Bool -> Grapher () -> Grapher ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (Operands -> Bool
    IS.null Operands
    ops) ((Binding -> Grapher ()) -> [Binding] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ Binding -> Grapher ()
    addVertex [Binding]
    bs Grapher () -> Grapher () -> Grapher ()
    forall a b.
    StateT State (Reader Env) a
    -> StateT State (Reader Env) b -> StateT State (Reader Env) b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Edges -> Operands -> Grapher ()
    addEdges Edges
    edges Operands
    ops)

    -- | Graph a statement that reads a scalar from device memory.
    graphRead :: Binding -> Grapher ()
    graphRead :: Binding -> Grapher ()
    graphRead Binding
    b = do
      -- Operands are not important as the source will block routes through b.
      Binding -> Grapher ()
    addSource Binding
    b
      Grapher ()
    tellRead

    -- | Graph a statement that always should be moved to device.
    graphAutoMove :: Binding -> Grapher ()
    graphAutoMove :: Binding -> Grapher ()
    graphAutoMove =
      -- Operands are not important as the source will block routes through b.
      Binding -> Grapher ()
    addSource

    -- | Graph a statement that is unfit for execution in a GPUBody and thus must
    -- be executed on host, requiring all its operands to be made available there.
    -- Parent statements of enclosing bodies are also blocked from being migrated.
    graphHostOnly :: Exp GPU -> Grapher ()
    graphHostOnly :: Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e = do
      -- Connect the vertices of all operands to sinks to mark that they are
      -- required on host. Transitive reads that they depend upon can be delayed
      -- no further, and any parent statements cannot be migrated.
      Operands
    ops <- Exp GPU -> Grapher Operands
    graphedScalarOperands Exp GPU
    e
      Edges -> Operands -> Grapher ()
    addEdges Edges
    ToSink Operands
    ops
      Grapher ()
    tellHostOnly

    -- | Graph an 'UpdateAcc' statement.
    graphUpdateAcc :: Binding -> Exp GPU -> Grapher ()
    graphUpdateAcc :: Binding -> Exp GPU -> Grapher ()
    graphUpdateAcc Binding
    b Exp GPU
    e | (Id
    _, Acc VName
    a ShapeBase SubExp
    _ [Type]
    _ NoUniqueness
    _) <- Binding
    b =
      -- The actual graphing is delayed to the corrensponding 'WithAcc' parent.
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st ->
        let accs :: IntMap [Delayed]
    accs = State -> IntMap [Delayed]
    stateUpdateAccs State
    st
            accs' :: IntMap [Delayed]
    accs' = (Maybe [Delayed] -> Maybe [Delayed])
    -> Id -> IntMap [Delayed] -> IntMap [Delayed]
    forall a. (Maybe a -> Maybe a) -> Id -> IntMap a -> IntMap a
    IM.alter Maybe [Delayed] -> Maybe [Delayed]
    add (VName -> Id
    nameToId VName
    a) IntMap [Delayed]
    accs
         in State
    st {stateUpdateAccs = accs'}
      where
        add :: Maybe [Delayed] -> Maybe [Delayed]
    add Maybe [Delayed]
    Nothing = [Delayed] -> Maybe [Delayed]
    forall a. a -> Maybe a
    Just [(Binding
    b, Exp GPU
    e)]
        add (Just [Delayed]
    xs) = [Delayed] -> Maybe [Delayed]
    forall a. a -> Maybe a
    Just ([Delayed] -> Maybe [Delayed]) -> [Delayed] -> Maybe [Delayed]
    forall a b. (a -> b) -> a -> b
    $ (Binding
    b, Exp GPU
    e) Delayed -> [Delayed] -> [Delayed]
    forall a. a -> [a] -> [a]
    : [Delayed]
    xs
    graphUpdateAcc Binding
    _ Exp GPU
    _ =
      String -> Grapher ()
    forall a. String -> a
    compilerBugS
        String
    "Type error: UpdateAcc did not produce accumulator typed value."

    -- | Graph a function application.
    graphApply :: Name -> [Binding] -> Exp GPU -> Grapher ()
    graphApply :: Name -> [Binding] -> Exp GPU -> Grapher ()
    graphApply Name
    fn [Binding]
    bs Exp GPU
    e = do
      Bool
    hof <- Name -> Grapher Bool
    isHostOnlyFun Name
    fn
      if Bool
    hof
        then Exp GPU -> Grapher ()
    graphHostOnly Exp GPU
    e
        else [Binding] -> Exp GPU -> Grapher ()
    graphSimple [Binding]
    bs Exp GPU
    e

    -- | Graph a Match statement.
    graphMatch :: [Binding] -> [SubExp] -> [Case (Body GPU)] -> Body GPU -> Grapher ()
    graphMatch :: [Binding]
    -> [SubExp] -> [Case (Body GPU)] -> Body GPU -> Grapher ()
    graphMatch [Binding]
    bs [SubExp]
    ses [Case (Body GPU)]
    cases Body GPU
    defbody = do
      Bool
    body_host_only <-
        Grapher Bool -> Grapher Bool
    forall a. Grapher a -> Grapher a
    incForkDepthFor (Grapher Bool -> Grapher Bool) -> Grapher Bool -> Grapher Bool
    forall a b. (a -> b) -> a -> b
    $
          (BodyStats -> Bool) -> [BodyStats] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any BodyStats -> Bool
    bodyHostOnly
            ([BodyStats] -> Bool)
    -> StateT State (Reader Env) [BodyStats] -> Grapher Bool
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Body GPU -> Grapher BodyStats)
    -> [Body GPU] -> StateT State (Reader Env) [BodyStats]
    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 (Grapher () -> Grapher BodyStats
    forall a. Grapher a -> Grapher BodyStats
    captureBodyStats (Grapher () -> Grapher BodyStats)
    -> (Body GPU -> Grapher ()) -> Body GPU -> Grapher BodyStats
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body GPU -> Grapher ()
    graphBody) (Body GPU
    defbody Body GPU -> [Body GPU] -> [Body GPU]
    forall a. a -> [a] -> [a]
    : (Case (Body GPU) -> Body GPU) -> [Case (Body GPU)] -> [Body GPU]
    forall a b. (a -> b) -> [a] -> [b]
    map Case (Body GPU) -> Body GPU
    forall body. Case body -> body
    caseBody [Case (Body GPU)]
    cases)

      let branch_results :: [[SubExp]]
    branch_results = Body GPU -> [SubExp]
    forall {rep}. Body rep -> [SubExp]
    results Body GPU
    defbody [SubExp] -> [[SubExp]] -> [[SubExp]]
    forall a. a -> [a] -> [a]
    : (Case (Body GPU) -> [SubExp]) -> [Case (Body GPU)] -> [[SubExp]]
    forall a b. (a -> b) -> [a] -> [b]
    map (Body GPU -> [SubExp]
    forall {rep}. Body rep -> [SubExp]
    results (Body GPU -> [SubExp])
    -> (Case (Body GPU) -> Body GPU) -> Case (Body GPU) -> [SubExp]
    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

      -- Record aliases for copyable memory backing returned arrays.
      Bool
    may_copy_results <- [Binding] -> [[SubExp]] -> Grapher Bool
    reusesBranches [Binding]
    bs [[SubExp]]
    branch_results
      let may_migrate :: Bool
    may_migrate = Bool -> Bool
    not Bool
    body_host_only Bool -> Bool -> Bool
    && Bool
    may_copy_results

      Operands
    cond_id <-
        if Bool
    may_migrate
          then [VName] -> Grapher Operands
    forall (t :: * -> *). Foldable t => t VName -> Grapher Operands
    onlyGraphedScalars ([VName] -> Grapher Operands) -> [VName] -> Grapher Operands
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> [VName]
    subExpVars [SubExp]
    ses
          else do
            -- The migration status of the condition is what determines
            -- whether the statement may be migrated as a whole or
            -- not. See 'shouldMoveStm'.
            (VName -> Grapher ()) -> [VName] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Id -> Grapher ()
    connectToSink (Id -> Grapher ()) -> (VName -> Id) -> VName -> Grapher ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> Id
    nameToId) ([SubExp] -> [VName]
    subExpVars [SubExp]
    ses)
            Operands -> Grapher Operands
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Operands
    IS.empty

      Operands -> Grapher ()
    tellOperands Operands
    cond_id

      -- Connect branch results to bound variables to allow delaying reads out of
      -- branches. It might also be beneficial to move the whole statement to
      -- device, to avoid reading the branch condition value. This must be balanced
      -- against the need to read the values bound by the if statement.
      --
      -- By connecting the branch condition to each variable bound by the statement
      -- the condition will only stay on device if
      --
      --   (1) the if statement is not required on host, based on the statements
      --       within its body.
      --
      --   (2) no additional reads will be required to use the if statement bound
      --       variables should the whole statement be migrated.
      --
      -- If the condition is migrated to device and stays there, then the if
      -- statement must necessarily execute on device.
      --
      -- While the graph model built by this module generally migrates no more
      -- statements than necessary to obtain a minimum vertex cut, the branches
      -- of if statements are subject to an inaccuracy. Specifically model is not
      -- strong enough to capture their mutual exclusivity and thus encodes that
      -- both branches are taken. While this does not affect the resulting number
      -- of host-device reads it means that some reads may needlessly be delayed
      -- out of branches. The overhead as measured on futhark-benchmarks appears
      -- to be neglible though.
      [Operands]
    ret <- ([SubExp] -> Grapher Operands)
    -> [[SubExp]] -> StateT State (Reader Env) [Operands]
    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 (Operands -> [SubExp] -> Grapher Operands
    comb Operands
    cond_id) ([[SubExp]] -> StateT State (Reader Env) [Operands])
    -> [[SubExp]] -> StateT State (Reader Env) [Operands]
    forall a b. (a -> b) -> a -> b
    $ [[SubExp]] -> [[SubExp]]
    forall a. [[a]] -> [[a]]
    L.transpose [[SubExp]]
    branch_results
      ((Binding, Operands) -> Grapher ())
    -> [(Binding, Operands)] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ ((Binding -> Operands -> Grapher ())
    -> (Binding, Operands) -> Grapher ()
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry Binding -> Operands -> Grapher ()
    createNode) ([Binding] -> [Operands] -> [(Binding, Operands)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Binding]
    bs [Operands]
    ret)
      where
        results :: Body rep -> [SubExp]
    results = (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp ([SubExpRes] -> [SubExp])
    -> (Body rep -> [SubExpRes]) -> Body rep -> [SubExp]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body rep -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult

        comb :: Operands -> [SubExp] -> Grapher Operands
    comb Operands
    ci [SubExp]
    a = (Operands
    ci <>) (Operands -> Operands) -> Grapher Operands -> Grapher Operands
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Set VName -> Grapher Operands
    forall (t :: * -> *). Foldable t => t VName -> Grapher Operands
    onlyGraphedScalars ([VName] -> Set VName
    forall a. Ord a => [a] -> Set a
    S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> [VName]
    subExpVars [SubExp]
    a)

    -----------------------------------------------------
    -- These type aliases are only used by 'graphLoop' --
    -----------------------------------------------------
    type ReachableBindings = IdSet

    type ReachableBindingsCache = MG.Visited (MG.Result ReachableBindings)

    type NonExhausted = [Id]

    type LoopValue = (Binding, Id, SubExp, SubExp)

    -----------------------------------------------------
    -----------------------------------------------------

    -- | Graph a loop statement.
    graphLoop ::
      [Binding] ->
      [(FParam GPU, SubExp)] ->
      LoopForm ->
      Body GPU ->
      Grapher ()
    graphLoop :: [Binding]
    -> [(FParam GPU, SubExp)] -> LoopForm -> Body GPU -> Grapher ()
    graphLoop [] [(FParam GPU, SubExp)]
    _ LoopForm
    _ Body GPU
    _ =
      -- We expect each loop to bind a value or be eliminated.
      String -> Grapher ()
    forall a. String -> a
    compilerBugS String
    "Loop statement bound no variable; should have been eliminated."
    graphLoop (Binding
    b : [Binding]
    bs) [(FParam GPU, SubExp)]
    params LoopForm
    lform Body GPU
    body = do
      -- Graph loop params and body while capturing statistics.
      Graph
    g <- Grapher Graph
    getGraph
      BodyStats
    stats <- Grapher () -> Grapher BodyStats
    forall a. Grapher a -> Grapher BodyStats
    captureBodyStats (Id
    subgraphId Id -> Grapher () -> Grapher ()
    forall a. Id -> Grapher a -> Grapher a
    `graphIdFor` Grapher ()
    graphTheLoop)

      -- Record aliases for copyable memory backing returned arrays.
      -- Does the loop return any arrays which prevent it from being migrated?
      let args :: [SubExp]
    args = ((Param DeclType, SubExp) -> SubExp)
    -> [(Param DeclType, SubExp)] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param DeclType, SubExp) -> SubExp
    forall a b. (a, b) -> b
    snd [(Param DeclType, SubExp)]
    [(FParam GPU, SubExp)]
    params
      let results :: [SubExp]
    results = (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp (Body GPU -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult Body GPU
    body)
      Bool
    may_copy_results <- [Binding] -> [[SubExp]] -> Grapher Bool
    reusesBranches (Binding
    b Binding -> [Binding] -> [Binding]
    forall a. a -> [a] -> [a]
    : [Binding]
    bs) [[SubExp]
    args, [SubExp]
    results]

      -- Connect the loop condition to a sink if the loop cannot be migrated,
      -- ensuring that it will be available to the host. The migration status
      -- of the condition is what determines whether the loop may be migrated
      -- as a whole or not. See 'shouldMoveStm'.
      let may_migrate :: Bool
    may_migrate = Bool -> Bool
    not (BodyStats -> Bool
    bodyHostOnly BodyStats
    stats) Bool -> Bool -> Bool
    && Bool
    may_copy_results
      Bool -> Grapher () -> Grapher ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless Bool
    may_migrate (Grapher () -> Grapher ()) -> Grapher () -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ case LoopForm
    lform of
        ForLoop VName
    _ IntType
    _ (Var VName
    n) -> Id -> Grapher ()
    connectToSink (VName -> Id
    nameToId VName
    n)
        WhileLoop VName
    n
          | Just (Binding
    _, Id
    p, SubExp
    _, SubExp
    res) <- VName -> Maybe (Binding, Id, SubExp, SubExp)
    loopValueFor VName
    n -> do
              Id -> Grapher ()
    connectToSink Id
    p
              case SubExp
    res of
                Var VName
    v -> Id -> Grapher ()
    connectToSink (VName -> Id
    nameToId VName
    v)
                SubExp
    _ -> () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
        LoopForm
    _ -> () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

      -- Connect graphed return values to their loop parameters.
      ((Binding, Id, SubExp, SubExp) -> Grapher ())
    -> [(Binding, Id, SubExp, SubExp)] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Binding, Id, SubExp, SubExp) -> Grapher ()
    mergeLoopParam [(Binding, Id, SubExp, SubExp)]
    loopValues

      -- Route the sources within the loop body in isolation.
      -- The loop graph must not be altered after this point.
      [Id]
    srcs <- Id -> Grapher [Id]
    routeSubgraph Id
    subgraphId

      -- Graph the variables bound by the statement.
      [(Binding, Id, SubExp, SubExp)]
    -> ((Binding, Id, SubExp, SubExp) -> Grapher ()) -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [(Binding, Id, SubExp, SubExp)]
    loopValues (((Binding, Id, SubExp, SubExp) -> Grapher ()) -> Grapher ())
    -> ((Binding, Id, SubExp, SubExp) -> Grapher ()) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \(Binding
    bnd, Id
    p, SubExp
    _, SubExp
    _) -> Binding -> Operands -> Grapher ()
    createNode Binding
    bnd (Id -> Operands
    IS.singleton Id
    p)

      -- If a device read is delayed from one iteration to the next the
      -- corresponding variables bound by the statement must be treated as
      -- sources.
      Graph
    g' <- Grapher Graph
    getGraph
      let (Operands
    dbs, ReachableBindingsCache
    rbc) = ((Operands, ReachableBindingsCache)
     -> Id -> (Operands, ReachableBindingsCache))
    -> (Operands, ReachableBindingsCache)
    -> [Id]
    -> (Operands, ReachableBindingsCache)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' (Graph
    -> (Operands, ReachableBindingsCache)
    -> Id
    -> (Operands, ReachableBindingsCache)
    deviceBindings Graph
    g') (Operands
    IS.empty, ReachableBindingsCache
    forall a. Visited a
    MG.none) [Id]
    srcs
      (Sources -> Sources) -> Grapher ()
    modifySources ((Sources -> Sources) -> Grapher ())
    -> (Sources -> Sources) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ ([Id] -> [Id]) -> Sources -> Sources
    forall b c a. (b -> c) -> (a, b) -> (a, c)
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second (Operands -> [Id]
    IS.toList Operands
    dbs <>)

      -- Connect operands to sinks if they can reach a sink within the loop.
      -- Otherwise connect them to the loop bound variables that they can
      -- reach and exhaust their normal entry edges into the loop.
      -- This means a read can be delayed through a loop but not into it if
      -- that would increase the number of reads done by any given iteration.
      let ops :: Operands
    ops = (Id -> Bool) -> Operands -> Operands
    IS.filter (Id -> Graph -> Bool
    forall m. Id -> Graph m -> Bool
    `MG.member` Graph
    g) (BodyStats -> Operands
    bodyOperands BodyStats
    stats)
      (ReachableBindingsCache
     -> Id -> StateT State (Reader Env) ReachableBindingsCache)
    -> ReachableBindingsCache -> [Id] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m ()
    foldM_ ReachableBindingsCache
    -> Id -> StateT State (Reader Env) ReachableBindingsCache
    connectOperand ReachableBindingsCache
    rbc (Operands -> [Id]
    IS.elems Operands
    ops)

      -- It might be beneficial to move the whole loop to device, to avoid
      -- reading the (initial) loop condition value. This must be balanced
      -- against the need to read the values bound by the loop statement.
      --
      -- For more details see the similar description for if statements.
      Bool -> Grapher () -> Grapher ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when Bool
    may_migrate (Grapher () -> Grapher ()) -> Grapher () -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ case LoopForm
    lform of
        ForLoop VName
    _ IntType
    _ SubExp
    n ->
          SubExp -> Grapher Operands
    onlyGraphedScalarSubExp SubExp
    n Grapher Operands -> (Operands -> Grapher ()) -> Grapher ()
    forall a b.
    StateT State (Reader Env) a
    -> (a -> StateT State (Reader Env) b)
    -> StateT State (Reader Env) b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Edges -> Operands -> Grapher ()
    addEdges (Operands -> Maybe Operands -> Edges
    ToNodes Operands
    bindings Maybe Operands
    forall a. Maybe a
    Nothing)
        WhileLoop VName
    n
          | Just (Binding
    _, Id
    _, SubExp
    arg, SubExp
    _) <- VName -> Maybe (Binding, Id, SubExp, SubExp)
    loopValueFor VName
    n ->
              SubExp -> Grapher Operands
    onlyGraphedScalarSubExp SubExp
    arg Grapher Operands -> (Operands -> Grapher ()) -> Grapher ()
    forall a b.
    StateT State (Reader Env) a
    -> (a -> StateT State (Reader Env) b)
    -> StateT State (Reader Env) b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Edges -> Operands -> Grapher ()
    addEdges (Operands -> Maybe Operands -> Edges
    ToNodes Operands
    bindings Maybe Operands
    forall a. Maybe a
    Nothing)
        LoopForm
    _ -> () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
      where
        subgraphId :: Id
        subgraphId :: Id
    subgraphId = Binding -> Id
    forall a b. (a, b) -> a
    fst Binding
    b

        loopValues :: [LoopValue]
        loopValues :: [(Binding, Id, SubExp, SubExp)]
    loopValues =
          let tmp :: [(Binding, (Param DeclType, SubExp), SubExpRes)]
    tmp = [Binding]
    -> [(Param DeclType, SubExp)]
    -> [SubExpRes]
    -> [(Binding, (Param DeclType, SubExp), SubExpRes)]
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
    zip3 (Binding
    b Binding -> [Binding] -> [Binding]
    forall a. a -> [a] -> [a]
    : [Binding]
    bs) [(Param DeclType, SubExp)]
    [(FParam GPU, SubExp)]
    params (Body GPU -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult Body GPU
    body)
              tmp' :: [(Binding, Id, SubExp, SubExp)]
    tmp' = (((Binding, (Param DeclType, SubExp), SubExpRes)
      -> (Binding, Id, SubExp, SubExp))
     -> [(Binding, (Param DeclType, SubExp), SubExpRes)]
     -> [(Binding, Id, SubExp, SubExp)])
    -> [(Binding, (Param DeclType, SubExp), SubExpRes)]
    -> ((Binding, (Param DeclType, SubExp), SubExpRes)
        -> (Binding, Id, SubExp, SubExp))
    -> [(Binding, Id, SubExp, SubExp)]
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip ((Binding, (Param DeclType, SubExp), SubExpRes)
     -> (Binding, Id, SubExp, SubExp))
    -> [(Binding, (Param DeclType, SubExp), SubExpRes)]
    -> [(Binding, Id, SubExp, SubExp)]
    forall a b. (a -> b) -> [a] -> [b]
    map [(Binding, (Param DeclType, SubExp), SubExpRes)]
    tmp (((Binding, (Param DeclType, SubExp), SubExpRes)
      -> (Binding, Id, SubExp, SubExp))
     -> [(Binding, Id, SubExp, SubExp)])
    -> ((Binding, (Param DeclType, SubExp), SubExpRes)
        -> (Binding, Id, SubExp, SubExp))
    -> [(Binding, Id, SubExp, SubExp)]
    forall a b. (a -> b) -> a -> b
    $
                \(Binding
    bnd, (Param DeclType
    p, SubExp
    arg), SubExpRes
    res) ->
                  let i :: Id
    i = VName -> Id
    nameToId (Param DeclType -> VName
    forall dec. Param dec -> VName
    paramName Param DeclType
    p)
                   in (Binding
    bnd, Id
    i, SubExp
    arg, SubExpRes -> SubExp
    resSubExp SubExpRes
    res)
           in ((Binding, Id, SubExp, SubExp) -> Bool)
    -> [(Binding, Id, SubExp, SubExp)]
    -> [(Binding, Id, SubExp, SubExp)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (\((Id
    _, Type
    t), Id
    _, SubExp
    _, SubExp
    _) -> Type -> Bool
    forall t. Typed t => t -> Bool
    isScalar Type
    t) [(Binding, Id, SubExp, SubExp)]
    tmp'

        bindings :: IdSet
        bindings :: Operands
    bindings = [Id] -> Operands
    IS.fromList ([Id] -> Operands) -> [Id] -> Operands
    forall a b. (a -> b) -> a -> b
    $ ((Binding, Id, SubExp, SubExp) -> Id)
    -> [(Binding, Id, SubExp, SubExp)] -> [Id]
    forall a b. (a -> b) -> [a] -> [b]
    map (\((Id
    i, Type
    _), Id
    _, SubExp
    _, SubExp
    _) -> Id
    i) [(Binding, Id, SubExp, SubExp)]
    loopValues

        loopValueFor :: VName -> Maybe (Binding, Id, SubExp, SubExp)
    loopValueFor VName
    n =
          ((Binding, Id, SubExp, SubExp) -> Bool)
    -> [(Binding, Id, SubExp, SubExp)]
    -> Maybe (Binding, Id, SubExp, SubExp)
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
    find (\(Binding
    _, Id
    p, SubExp
    _, SubExp
    _) -> Id
    p Id -> Id -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName -> Id
    nameToId VName
    n) [(Binding, Id, SubExp, SubExp)]
    loopValues

        graphTheLoop :: Grapher ()
        graphTheLoop :: Grapher ()
    graphTheLoop = do
          ((Binding, Id, SubExp, SubExp) -> Grapher ())
    -> [(Binding, Id, SubExp, SubExp)] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Binding, Id, SubExp, SubExp) -> Grapher ()
    forall {a} {d}. ((a, Type), Id, SubExp, d) -> Grapher ()
    graphParam [(Binding, Id, SubExp, SubExp)]
    loopValues

          -- For simplicity we do not currently track memory reuse through merge
          -- parameters. A parameter does not simply reuse the memory of its
          -- argument; it must also consider the iteration return value, which in
          -- turn may depend on other merge parameters.
          --
          -- Situations that would benefit from this tracking is unlikely to occur
          -- at the time of writing, and if it occurs current compiler limitations
          -- will prevent successful compilation.
          -- Specifically it requires the merge parameter argument to reuse memory
          -- from an array literal, and both it and the loop must occur within an
          -- if statement branch. Array literals are generally hoisted out of if
          -- statements however, and when they are not, a memory allocation error
          -- occurs.
          --
          -- TODO: Track memory reuse through merge parameters.

          case LoopForm
    lform of
            ForLoop VName
    _ IntType
    _ SubExp
    n ->
              SubExp -> Grapher Operands
    onlyGraphedScalarSubExp SubExp
    n Grapher Operands -> (Operands -> Grapher ()) -> Grapher ()
    forall a b.
    StateT State (Reader Env) a
    -> (a -> StateT State (Reader Env) b)
    -> StateT State (Reader Env) b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Operands -> Grapher ()
    tellOperands
            WhileLoop VName
    _ -> () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
          Body GPU -> Grapher ()
    graphBody Body GPU
    body
          where
            graphParam :: ((a, Type), Id, SubExp, d) -> Grapher ()
    graphParam ((a
    _, Type
    t), Id
    p, SubExp
    arg, d
    _) =
              do
                -- It is unknown whether a read can be delayed via the parameter
                -- from one iteration to the next, so we have to create a vertex
                -- even if the initial value never depends on a read.
                Binding -> Grapher ()
    addVertex (Id
    p, Type
    t)
                Operands
    ops <- SubExp -> Grapher Operands
    onlyGraphedScalarSubExp SubExp
    arg
                Edges -> Operands -> Grapher ()
    addEdges (Id -> Edges
    MG.oneEdge Id
    p) Operands
    ops

        mergeLoopParam :: LoopValue -> Grapher ()
        mergeLoopParam :: (Binding, Id, SubExp, SubExp) -> Grapher ()
    mergeLoopParam (Binding
    _, Id
    p, SubExp
    _, SubExp
    res)
          | Var VName
    n <- SubExp
    res,
            Id
    ret <- VName -> Id
    nameToId VName
    n,
            Id
    ret Id -> Id -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Id
    p =
              Edges -> Operands -> Grapher ()
    addEdges (Id -> Edges
    MG.oneEdge Id
    p) (Id -> Operands
    IS.singleton Id
    ret)
          | Bool
    otherwise =
              () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
        deviceBindings ::
          Graph ->
          (ReachableBindings, ReachableBindingsCache) ->
          Id ->
          (ReachableBindings, ReachableBindingsCache)
        deviceBindings :: Graph
    -> (Operands, ReachableBindingsCache)
    -> Id
    -> (Operands, ReachableBindingsCache)
    deviceBindings Graph
    g (Operands
    rb, ReachableBindingsCache
    rbc) Id
    i =
          let (Result Operands
    r, ReachableBindingsCache
    rbc') = Graph
    -> (Operands -> EdgeType -> Vertex Meta -> Operands)
    -> ReachableBindingsCache
    -> EdgeType
    -> Id
    -> (Result Operands, ReachableBindingsCache)
    forall a m.
    Monoid a =>
    Graph m
    -> (a -> EdgeType -> Vertex m -> a)
    -> Visited (Result a)
    -> EdgeType
    -> Id
    -> (Result a, Visited (Result a))
    MG.reduce Graph
    g Operands -> EdgeType -> Vertex Meta -> Operands
    bindingReach ReachableBindingsCache
    rbc EdgeType
    Normal Id
    i
           in case Result Operands
    r of
                Produced Operands
    rb' -> (Operands
    rb Operands -> Operands -> Operands
    forall a. Semigroup a => a -> a -> a
    <> Operands
    rb', ReachableBindingsCache
    rbc')
                Result Operands
    _ ->
                  String -> (Operands, ReachableBindingsCache)
    forall a. String -> a
    compilerBugS
                    String
    "Migration graph sink could be reached from source after it\
                    \ had been attempted routed."
        bindingReach ::
          ReachableBindings ->
          EdgeType ->
          Vertex Meta ->
          ReachableBindings
        bindingReach :: Operands -> EdgeType -> Vertex Meta -> Operands
    bindingReach Operands
    rb EdgeType
    _ Vertex Meta
    v
          | Id
    i <- Vertex Meta -> Id
    forall m. Vertex m -> Id
    vertexId Vertex Meta
    v,
            Id -> Operands -> Bool
    IS.member Id
    i Operands
    bindings =
              Id -> Operands -> Operands
    IS.insert Id
    i Operands
    rb
          | Bool
    otherwise =
              Operands
    rb
        connectOperand ::
          ReachableBindingsCache ->
          Id ->
          Grapher ReachableBindingsCache
        connectOperand :: ReachableBindingsCache
    -> Id -> StateT State (Reader Env) ReachableBindingsCache
    connectOperand ReachableBindingsCache
    cache Id
    op = do
          Graph
    g <- Grapher Graph
    getGraph
          case Id -> Graph -> Maybe (Vertex Meta)
    forall m. Id -> Graph m -> Maybe (Vertex m)
    MG.lookup Id
    op Graph
    g of
            Maybe (Vertex Meta)
    Nothing -> ReachableBindingsCache
    -> StateT State (Reader Env) ReachableBindingsCache
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ReachableBindingsCache
    cache
            Just Vertex Meta
    v ->
              case Vertex Meta -> Edges
    forall m. Vertex m -> Edges
    vertexEdges Vertex Meta
    v of
                Edges
    ToSink -> ReachableBindingsCache
    -> StateT State (Reader Env) ReachableBindingsCache
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ReachableBindingsCache
    cache
                ToNodes Operands
    es Maybe Operands
    Nothing -> Graph
    -> ReachableBindingsCache
    -> Id
    -> Operands
    -> StateT State (Reader Env) ReachableBindingsCache
    connectOp Graph
    g ReachableBindingsCache
    cache Id
    op Operands
    es
                ToNodes Operands
    _ (Just Operands
    nx) -> Graph
    -> ReachableBindingsCache
    -> Id
    -> Operands
    -> StateT State (Reader Env) ReachableBindingsCache
    connectOp Graph
    g ReachableBindingsCache
    cache Id
    op Operands
    nx
          where
            connectOp ::
              Graph ->
              ReachableBindingsCache ->
              Id -> -- operand id
              IdSet -> -- its edges
              Grapher ReachableBindingsCache
            connectOp :: Graph
    -> ReachableBindingsCache
    -> Id
    -> Operands
    -> StateT State (Reader Env) ReachableBindingsCache
    connectOp Graph
    g ReachableBindingsCache
    rbc Id
    i Operands
    es = do
              let (Result Operands
    res, [Id]
    nx, ReachableBindingsCache
    rbc') = Graph
    -> (Operands, [Id], ReachableBindingsCache)
    -> [Id]
    -> (Result Operands, [Id], ReachableBindingsCache)
    findBindings Graph
    g (Operands
    IS.empty, [], ReachableBindingsCache
    rbc) (Operands -> [Id]
    IS.elems Operands
    es)
              case Result Operands
    res of
                Result Operands
    FoundSink -> Id -> Grapher ()
    connectToSink Id
    i
                Produced Operands
    rb -> (Graph -> Graph) -> Grapher ()
    modifyGraph ((Graph -> Graph) -> Grapher ()) -> (Graph -> Graph) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ (Vertex Meta -> Vertex Meta) -> Id -> Graph -> Graph
    forall m. (Vertex m -> Vertex m) -> Id -> Graph m -> Graph m
    MG.adjust ([Id] -> Operands -> Vertex Meta -> Vertex Meta
    updateEdges [Id]
    nx Operands
    rb) Id
    i
              ReachableBindingsCache
    -> StateT State (Reader Env) ReachableBindingsCache
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ReachableBindingsCache
    rbc'
            updateEdges ::
              NonExhausted ->
              ReachableBindings ->
              Vertex Meta ->
              Vertex Meta
            updateEdges :: [Id] -> Operands -> Vertex Meta -> Vertex Meta
    updateEdges [Id]
    nx Operands
    rb Vertex Meta
    v
              | ToNodes Operands
    es Maybe Operands
    _ <- Vertex Meta -> Edges
    forall m. Vertex m -> Edges
    vertexEdges Vertex Meta
    v =
                  let nx' :: Operands
    nx' = [Id] -> Operands
    IS.fromList [Id]
    nx
                      es' :: Edges
    es' = Operands -> Maybe Operands -> Edges
    ToNodes (Operands
    rb Operands -> Operands -> Operands
    forall a. Semigroup a => a -> a -> a
    <> Operands
    es) (Maybe Operands -> Edges) -> Maybe Operands -> Edges
    forall a b. (a -> b) -> a -> b
    $ Operands -> Maybe Operands
    forall a. a -> Maybe a
    Just (Operands
    rb Operands -> Operands -> Operands
    forall a. Semigroup a => a -> a -> a
    <> Operands
    nx')
                   in Vertex Meta
    v {vertexEdges = es'}
              | Bool
    otherwise = Vertex Meta
    v
            findBindings ::
              Graph ->
              (ReachableBindings, NonExhausted, ReachableBindingsCache) ->
              [Id] -> -- current non-exhausted edges
              (MG.Result ReachableBindings, NonExhausted, ReachableBindingsCache)
            findBindings :: Graph
    -> (Operands, [Id], ReachableBindingsCache)
    -> [Id]
    -> (Result Operands, [Id], ReachableBindingsCache)
    findBindings Graph
    _ (Operands
    rb, [Id]
    nx, ReachableBindingsCache
    rbc) [] =
              (Operands -> Result Operands
    forall a. a -> Result a
    Produced Operands
    rb, [Id]
    nx, ReachableBindingsCache
    rbc)
            findBindings Graph
    g (Operands
    rb, [Id]
    nx, ReachableBindingsCache
    rbc) (Id
    i : [Id]
    is)
              | Just Vertex Meta
    v <- Id -> Graph -> Maybe (Vertex Meta)
    forall m. Id -> Graph m -> Maybe (Vertex m)
    MG.lookup Id
    i Graph
    g,
                Just Id
    gid <- Meta -> Maybe Id
    metaGraphId (Vertex Meta -> Meta
    forall m. Vertex m -> m
    vertexMeta Vertex Meta
    v),
                Id
    gid Id -> Id -> Bool
    forall a. Eq a => a -> a -> Bool
    == Id
    subgraphId -- only search the subgraph
                =
                  let (Result Operands
    res, ReachableBindingsCache
    rbc') = Graph
    -> (Operands -> EdgeType -> Vertex Meta -> Operands)
    -> ReachableBindingsCache
    -> EdgeType
    -> Id
    -> (Result Operands, ReachableBindingsCache)
    forall a m.
    Monoid a =>
    Graph m
    -> (a -> EdgeType -> Vertex m -> a)
    -> Visited (Result a)
    -> EdgeType
    -> Id
    -> (Result a, Visited (Result a))
    MG.reduce Graph
    g Operands -> EdgeType -> Vertex Meta -> Operands
    bindingReach ReachableBindingsCache
    rbc EdgeType
    Normal Id
    i
                   in case Result Operands
    res of
                        Result Operands
    FoundSink -> (Result Operands
    forall a. Result a
    FoundSink, [], ReachableBindingsCache
    rbc')
                        Produced Operands
    rb' -> Graph
    -> (Operands, [Id], ReachableBindingsCache)
    -> [Id]
    -> (Result Operands, [Id], ReachableBindingsCache)
    findBindings Graph
    g (Operands
    rb Operands -> Operands -> Operands
    forall a. Semigroup a => a -> a -> a
    <> Operands
    rb', [Id]
    nx, ReachableBindingsCache
    rbc') [Id]
    is
              | Bool
    otherwise =
                  -- don't exhaust
                  Graph
    -> (Operands, [Id], ReachableBindingsCache)
    -> [Id]
    -> (Result Operands, [Id], ReachableBindingsCache)
    findBindings Graph
    g (Operands
    rb, Id
    i Id -> [Id] -> [Id]
    forall a. a -> [a] -> [a]
    : [Id]
    nx, ReachableBindingsCache
    rbc) [Id]
    is

    -- | Graph a 'WithAcc' statement.
    graphWithAcc ::
      [Binding] ->
      [WithAccInput GPU] ->
      Lambda GPU ->
      Grapher ()
    graphWithAcc :: [Binding] -> [WithAccInput GPU] -> Lambda GPU -> Grapher ()
    graphWithAcc [Binding]
    bs [WithAccInput GPU]
    inputs Lambda GPU
    f = do
      -- Graph the body, capturing 'UpdateAcc' statements for delayed graphing.
      Body GPU -> Grapher ()
    graphBody (Lambda GPU -> Body GPU
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPU
    f)

      -- Graph each accumulator monoid and its associated 'UpdateAcc' statements.
      ((Type, WithAccInput GPU) -> Grapher ())
    -> [(Type, WithAccInput GPU)] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Type, WithAccInput GPU) -> Grapher ()
    forall {shape} {u} {a} {b}.
    (TypeBase shape u, (a, b, Maybe (Lambda GPU, [SubExp])))
    -> Grapher ()
    graph ([(Type, WithAccInput GPU)] -> Grapher ())
    -> [(Type, WithAccInput GPU)] -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ [Type] -> [WithAccInput GPU] -> [(Type, WithAccInput GPU)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Lambda GPU -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda GPU
    f) [WithAccInput GPU]
    inputs

      -- Record aliases for the backing memory of each returned array.
      -- 'WithAcc' statements are never migrated as a whole and always returns
      -- arrays backed by memory allocated elsewhere.
      let arrs :: [SubExp]
    arrs = (WithAccInput GPU -> [SubExp]) -> [WithAccInput GPU] -> [SubExp]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (\(ShapeBase SubExp
    _, [VName]
    as, Maybe (Lambda GPU, [SubExp])
    _) -> (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    Var [VName]
    as) [WithAccInput GPU]
    inputs
      let res :: [SubExpRes]
    res = Id -> [SubExpRes] -> [SubExpRes]
    forall a. Id -> [a] -> [a]
    drop ([WithAccInput GPU] -> Id
    forall a. [a] -> Id
    forall (t :: * -> *) a. Foldable t => t a -> Id
    length [WithAccInput GPU]
    inputs) (Body GPU -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult (Body GPU -> [SubExpRes]) -> Body GPU -> [SubExpRes]
    forall a b. (a -> b) -> a -> b
    $ Lambda GPU -> Body GPU
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPU
    f)
      Bool
    _ <- [Binding] -> [SubExp] -> Grapher Bool
    reusesReturn [Binding]
    bs ([SubExp]
    arrs [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ (SubExpRes -> SubExp) -> [SubExpRes] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp [SubExpRes]
    res)

      -- Connect return variables to bound values. No outgoing edge exists
      -- from an accumulator vertex so skip those. Note that accumulators do
      -- not map to returned arrays one-to-one but one-to-many.
      [Operands]
    ret <- (SubExpRes -> Grapher Operands)
    -> [SubExpRes] -> StateT State (Reader Env) [Operands]
    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 (SubExp -> Grapher Operands
    onlyGraphedScalarSubExp (SubExp -> Grapher Operands)
    -> (SubExpRes -> SubExp) -> SubExpRes -> Grapher Operands
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SubExpRes -> SubExp
    resSubExp) [SubExpRes]
    res
      ((Binding, Operands) -> Grapher ())
    -> [(Binding, Operands)] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ ((Binding -> Operands -> Grapher ())
    -> (Binding, Operands) -> Grapher ()
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry Binding -> Operands -> Grapher ()
    createNode) ([(Binding, Operands)] -> Grapher ())
    -> [(Binding, Operands)] -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ [Binding] -> [Operands] -> [(Binding, Operands)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Id -> [Binding] -> [Binding]
    forall a. Id -> [a] -> [a]
    drop ([SubExp] -> Id
    forall a. [a] -> Id
    forall (t :: * -> *) a. Foldable t => t a -> Id
    length [SubExp]
    arrs) [Binding]
    bs) [Operands]
    ret
      where
        graph :: (TypeBase shape u, (a, b, Maybe (Lambda GPU, [SubExp])))
    -> Grapher ()
    graph (Acc VName
    a ShapeBase SubExp
    _ [Type]
    types u
    _, (a
    _, b
    _, Maybe (Lambda GPU, [SubExp])
    comb)) = do
          let i :: Id
    i = VName -> Id
    nameToId VName
    a

          [Delayed]
    delayed <- [Delayed] -> Maybe [Delayed] -> [Delayed]
    forall a. a -> Maybe a -> a
    fromMaybe [] (Maybe [Delayed] -> [Delayed])
    -> StateT State (Reader Env) (Maybe [Delayed])
    -> StateT State (Reader Env) [Delayed]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (State -> Maybe [Delayed])
    -> StateT State (Reader Env) (Maybe [Delayed])
    forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
    gets (Id -> IntMap [Delayed] -> Maybe [Delayed]
    forall a. Id -> IntMap a -> Maybe a
    IM.lookup Id
    i (IntMap [Delayed] -> Maybe [Delayed])
    -> (State -> IntMap [Delayed]) -> State -> Maybe [Delayed]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . State -> IntMap [Delayed]
    stateUpdateAccs)
          (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateUpdateAccs = IM.delete i (stateUpdateAccs st)}

          Id -> [Type] -> Maybe (Lambda GPU) -> [Delayed] -> Grapher ()
    graphAcc Id
    i [Type]
    types ((Lambda GPU, [SubExp]) -> Lambda GPU
    forall a b. (a, b) -> a
    fst ((Lambda GPU, [SubExp]) -> Lambda GPU)
    -> Maybe (Lambda GPU, [SubExp]) -> Maybe (Lambda GPU)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Maybe (Lambda GPU, [SubExp])
    comb) [Delayed]
    delayed

          -- Neutral elements must always be made available on host for 'WithAcc'
          -- to type check.
          (SubExp -> Grapher ()) -> [SubExp] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SubExp -> Grapher ()
    connectSubExpToSink ([SubExp] -> Grapher ()) -> [SubExp] -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ [SubExp]
    -> ((Lambda GPU, [SubExp]) -> [SubExp])
    -> Maybe (Lambda GPU, [SubExp])
    -> [SubExp]
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe [] (Lambda GPU, [SubExp]) -> [SubExp]
    forall a b. (a, b) -> b
    snd Maybe (Lambda GPU, [SubExp])
    comb
        graph (TypeBase shape u, (a, b, Maybe (Lambda GPU, [SubExp])))
    _ =
          String -> Grapher ()
    forall a. String -> a
    compilerBugS String
    "Type error: WithAcc expression did not return accumulator."

    -- Graph the operator and all 'UpdateAcc' statements associated with an
    -- accumulator.
    --
    -- The arguments are the 'Id' for the accumulator token, the element types of
    -- the accumulator/operator, its combining function if any, and all associated
    -- 'UpdateAcc' statements outside kernels.
    graphAcc :: Id -> [Type] -> Maybe (Lambda GPU) -> [Delayed] -> Grapher ()
    graphAcc :: Id -> [Type] -> Maybe (Lambda GPU) -> [Delayed] -> Grapher ()
    graphAcc Id
    i [Type]
    _ Maybe (Lambda GPU)
    _ [] = Binding -> Grapher ()
    addSource (Id
    i, PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    Unit) -- Only used on device.
    graphAcc Id
    i [Type]
    types Maybe (Lambda GPU)
    op [Delayed]
    delayed = do
      -- Accumulators are intended for use within SegOps but in principle the AST
      -- allows their 'UpdateAcc's to be used outside a kernel. This case handles
      -- that unlikely situation.

      Env
    env <- Grapher Env
    ask
      State
    st <- StateT State (Reader Env) State
    forall (m :: * -> *) s. Monad m => StateT s m s
    get

      -- Collect statistics about the operator statements.
      let lambda :: Lambda GPU
    lambda = Lambda GPU -> Maybe (Lambda GPU) -> Lambda GPU
    forall a. a -> Maybe a -> a
    fromMaybe ([LParam GPU] -> [Type] -> Body GPU -> Lambda GPU
    forall rep. [LParam rep] -> [Type] -> Body rep -> Lambda rep
    Lambda [] [] (BodyDec GPU -> Stms GPU -> [SubExpRes] -> Body GPU
    forall rep. BodyDec rep -> Stms rep -> [SubExpRes] -> Body rep
    Body () Stms GPU
    forall a. Seq a
    SQ.empty [])) Maybe (Lambda GPU)
    op
      let m :: Grapher ()
    m = Body GPU -> Grapher ()
    graphBody (Lambda GPU -> Body GPU
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPU
    lambda)
      let stats :: BodyStats
    stats = Reader Env BodyStats -> Env -> BodyStats
    forall r a. Reader r a -> r -> a
    R.runReader (Grapher BodyStats -> State -> Reader Env BodyStats
    forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
    evalStateT (Grapher () -> Grapher BodyStats
    forall a. Grapher a -> Grapher BodyStats
    captureBodyStats Grapher ()
    m) State
    st) Env
    env
      -- We treat GPUBody kernels as host-only to not bother rewriting them inside
      -- operators and to simplify the analysis. They are unlikely to occur anyway.
      --
      -- NOTE: Performance may degrade if a GPUBody is replaced with its contents
      --       but the containing operator is used on host.
      let host_only :: Bool
    host_only = BodyStats -> Bool
    bodyHostOnly BodyStats
    stats Bool -> Bool -> Bool
    || BodyStats -> Bool
    bodyHasGPUBody BodyStats
    stats

      -- op operands are read from arrays and written back so if any of the operands
      -- are scalar then a read can be avoided by moving the UpdateAcc usages to
      -- device. If the op itself performs scalar reads its UpdateAcc usages should
      -- also be moved.
      let does_read :: Bool
    does_read = BodyStats -> Bool
    bodyReads BodyStats
    stats Bool -> Bool -> Bool
    || (Type -> Bool) -> [Type] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any Type -> Bool
    forall t. Typed t => t -> Bool
    isScalar [Type]
    types

      -- Determine which external variables the operator depends upon.
      -- 'bodyOperands' cannot be used as it might exclude operands that were
      -- connected to sinks within the body, so instead we create an artifical
      -- expression to capture graphed operands from.
      Operands
    ops <- Exp GPU -> Grapher Operands
    graphedScalarOperands ([WithAccInput GPU] -> Lambda GPU -> Exp GPU
    forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
    WithAcc [] Lambda GPU
    lambda)

      case (Bool
    host_only, Bool
    does_read) of
        (Bool
    True, Bool
    _) -> do
          -- If the operator cannot run well in a GPUBody then all non-kernel
          -- UpdateAcc statements are host-only. The current analysis is ignorant
          -- of what happens inside kernels so we must assume that the operator
          -- is used within a kernel, meaning that we cannot migrate its statements.
          --
          -- TODO: Improve analysis if UpdateAcc ever is used outside kernels.
          (Delayed -> Grapher ()) -> [Delayed] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Exp GPU -> Grapher ()
    graphHostOnly (Exp GPU -> Grapher ())
    -> (Delayed -> Exp GPU) -> Delayed -> Grapher ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Delayed -> Exp GPU
    forall a b. (a, b) -> b
    snd) [Delayed]
    delayed
          Edges -> Operands -> Grapher ()
    addEdges Edges
    ToSink Operands
    ops
        (Bool
    _, Bool
    True) -> do
          -- Migrate all accumulator usage to device to avoid reads and writes.
          (Delayed -> Grapher ()) -> [Delayed] -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Binding -> Grapher ()
    graphAutoMove (Binding -> Grapher ())
    -> (Delayed -> Binding) -> Delayed -> Grapher ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Delayed -> Binding
    forall a b. (a, b) -> a
    fst) [Delayed]
    delayed
          Binding -> Grapher ()
    addSource (Id
    i, PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    Unit)
        (Bool, Bool)
    _ -> do
          -- Only migrate operator and UpdateAcc statements if it can allow their
          -- operands to be migrated.
          Binding -> Operands -> Grapher ()
    createNode (Id
    i, PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    Unit) Operands
    ops
          [Delayed] -> (Delayed -> Grapher ()) -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [Delayed]
    delayed ((Delayed -> Grapher ()) -> Grapher ())
    -> (Delayed -> Grapher ()) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $
            \(Binding
    b, Exp GPU
    e) -> Exp GPU -> Grapher Operands
    graphedScalarOperands Exp GPU
    e Grapher Operands -> (Operands -> Grapher ()) -> Grapher ()
    forall a b.
    StateT State (Reader Env) a
    -> (a -> StateT State (Reader Env) b)
    -> StateT State (Reader Env) b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Binding -> Operands -> Grapher ()
    createNode Binding
    b (Operands -> Grapher ())
    -> (Operands -> Operands) -> Operands -> Grapher ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Id -> Operands -> Operands
    IS.insert Id
    i

    -- Returns for an expression all scalar operands that must be made available
    -- on host to execute the expression there.
    graphedScalarOperands :: Exp GPU -> Grapher Operands
    graphedScalarOperands :: Exp GPU -> Grapher Operands
    graphedScalarOperands Exp GPU
    e =
      let is :: Operands
    is = (Operands, Set VName) -> Operands
    forall a b. (a, b) -> a
    fst ((Operands, Set VName) -> Operands)
    -> (Operands, Set VName) -> Operands
    forall a b. (a -> b) -> a -> b
    $ State (Operands, Set VName) ()
    -> (Operands, Set VName) -> (Operands, Set VName)
    forall s a. State s a -> s -> s
    execState (Exp GPU -> State (Operands, Set VName) ()
    collect Exp GPU
    e) (Operands, Set VName)
    forall {a}. (Operands, Set a)
    initial
       in Operands -> Operands -> Operands
    IS.intersection Operands
    is (Operands -> Operands) -> Grapher Operands -> Grapher Operands
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Grapher Operands
    getGraphedScalars
      where
        initial :: (Operands, Set a)
    initial = (Operands
    IS.empty, Set a
    forall a. Set a
    S.empty) -- scalar operands, accumulator tokens
        captureName :: VName -> StateT (p Operands c) m ()
    captureName VName
    n = (p Operands c -> p Operands c) -> StateT (p Operands c) m ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((p Operands c -> p Operands c) -> StateT (p Operands c) m ())
    -> (p Operands c -> p Operands c) -> StateT (p Operands c) m ()
    forall a b. (a -> b) -> a -> b
    $ (Operands -> Operands) -> p Operands c -> p Operands c
    forall a b c. (a -> b) -> p a c -> p b c
    forall (p :: * -> * -> *) a b c.
    Bifunctor p =>
    (a -> b) -> p a c -> p b c
    first ((Operands -> Operands) -> p Operands c -> p Operands c)
    -> (Operands -> Operands) -> p Operands c -> p Operands c
    forall a b. (a -> b) -> a -> b
    $ Id -> Operands -> Operands
    IS.insert (VName -> Id
    nameToId VName
    n)
        captureAcc :: a -> StateT (p a (Set a)) m ()
    captureAcc a
    a = (p a (Set a) -> p a (Set a)) -> StateT (p a (Set a)) m ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((p a (Set a) -> p a (Set a)) -> StateT (p a (Set a)) m ())
    -> (p a (Set a) -> p a (Set a)) -> StateT (p a (Set a)) m ()
    forall a b. (a -> b) -> a -> b
    $ (Set a -> Set a) -> p a (Set a) -> p a (Set a)
    forall b c a. (b -> c) -> p a b -> p a c
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second ((Set a -> Set a) -> p a (Set a) -> p a (Set a))
    -> (Set a -> Set a) -> p a (Set a) -> p a (Set a)
    forall a b. (a -> b) -> a -> b
    $ a -> Set a -> Set a
    forall a. Ord a => a -> Set a -> Set a
    S.insert a
    a
        collectFree :: a -> StateT (p Operands c) m ()
    collectFree a
    x = (VName -> StateT (p Operands c) m ())
    -> [VName] -> StateT (p Operands c) m ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ VName -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    VName -> StateT (p Operands c) m ()
    captureName (Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ a -> Names
    forall a. FreeIn a => a -> Names
    freeIn a
    x)

        collect :: Exp GPU -> State (Operands, Set VName) ()
    collect b :: Exp GPU
    b@BasicOp {} =
          Exp GPU -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {rep} {c}.
    (Monad m, Bifunctor p) =>
    Exp rep -> StateT (p Operands c) m ()
    collectBasic Exp GPU
    b
        collect (Apply Name
    _ [(SubExp, Diet)]
    params [(RetType GPU, RetAls)]
    _ Safety
    _) =
          ((SubExp, Diet) -> State (Operands, Set VName) ())
    -> [(SubExp, Diet)] -> State (Operands, Set VName) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (SubExp -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp (SubExp -> State (Operands, Set VName) ())
    -> ((SubExp, Diet) -> SubExp)
    -> (SubExp, Diet)
    -> State (Operands, Set VName) ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (SubExp, Diet) -> SubExp
    forall a b. (a, b) -> a
    fst) [(SubExp, Diet)]
    params
        collect (Match [SubExp]
    ses [Case (Body GPU)]
    cases Body GPU
    defbody MatchDec (BranchType GPU)
    _) = do
          (SubExp -> State (Operands, Set VName) ())
    -> [SubExp] -> State (Operands, Set VName) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SubExp -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp [SubExp]
    ses
          (Case (Body GPU) -> State (Operands, Set VName) ())
    -> [Case (Body GPU)] -> State (Operands, Set VName) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Body GPU -> State (Operands, Set VName) ()
    collectBody (Body GPU -> State (Operands, Set VName) ())
    -> (Case (Body GPU) -> Body GPU)
    -> Case (Body GPU)
    -> State (Operands, Set VName) ()
    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
          Body GPU -> State (Operands, Set VName) ()
    collectBody Body GPU
    defbody
        collect (Loop [(FParam GPU, SubExp)]
    params LoopForm
    lform Body GPU
    body) = do
          ((FParam GPU, SubExp) -> State (Operands, Set VName) ())
    -> [(FParam GPU, SubExp)] -> State (Operands, Set VName) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (SubExp -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp (SubExp -> State (Operands, Set VName) ())
    -> ((FParam GPU, SubExp) -> SubExp)
    -> (FParam GPU, SubExp)
    -> State (Operands, Set VName) ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (FParam GPU, SubExp) -> SubExp
    forall a b. (a, b) -> b
    snd) [(FParam GPU, SubExp)]
    params
          LoopForm -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    LoopForm -> StateT (p Operands c) m ()
    collectLForm LoopForm
    lform
          Body GPU -> State (Operands, Set VName) ()
    collectBody Body GPU
    body
        collect (WithAcc [WithAccInput GPU]
    accs Lambda GPU
    f) =
          [WithAccInput GPU] -> Lambda GPU -> State (Operands, Set VName) ()
    collectWithAcc [WithAccInput GPU]
    accs Lambda GPU
    f
        collect (Op Op GPU
    op) =
          HostOp SOAC GPU -> State (Operands, Set VName) ()
    forall {op :: * -> *} {rep} {c}.
    FreeIn (op rep) =>
    HostOp op rep -> StateT (Operands, c) Identity ()
    collectHostOp Op GPU
    HostOp SOAC GPU
    op

        collectBasic :: Exp rep -> StateT (p Operands c) m ()
    collectBasic (BasicOp (Update Safety
    _ VName
    _ Slice SubExp
    slice SubExp
    _)) =
          -- Writing a scalar to an array can be replaced with copying a single-
          -- element slice. If the scalar originates from device memory its read
          -- can thus be prevented without requiring the 'Update' to be migrated.
          Slice SubExp -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {a} {c}.
    (Monad m, Bifunctor p, FreeIn a) =>
    a -> StateT (p Operands c) m ()
    collectFree Slice SubExp
    slice
        collectBasic (BasicOp (Replicate ShapeBase SubExp
    shape SubExp
    _)) =
          -- The replicate of a scalar can be rewritten as a replicate of a single
          -- element array followed by a slice index.
          ShapeBase SubExp -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {a} {c}.
    (Monad m, Bifunctor p, FreeIn a) =>
    a -> StateT (p Operands c) m ()
    collectFree ShapeBase SubExp
    shape
        collectBasic Exp rep
    e' =
          -- Note: Plain VName values only refer to arrays.
          Walker rep (StateT (p Operands c) m)
    -> Exp rep -> StateT (p Operands c) m ()
    forall (m :: * -> *) rep.
    Monad m =>
    Walker rep m -> Exp rep -> m ()
    walkExpM (Walker rep (StateT (p Operands c) m)
    forall rep (m :: * -> *). Monad m => Walker rep m
    identityWalker {walkOnSubExp = collectSubExp}) Exp rep
    e'

        collectSubExp :: SubExp -> StateT (p Operands c) m ()
    collectSubExp (Var VName
    n) = VName -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    VName -> StateT (p Operands c) m ()
    captureName VName
    n
        collectSubExp SubExp
    _ = () -> StateT (p Operands c) m ()
    forall a. a -> StateT (p Operands c) m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

        collectBody :: Body GPU -> State (Operands, Set VName) ()
    collectBody Body GPU
    body = do
          Stms GPU -> State (Operands, Set VName) ()
    collectStms (Body GPU -> Stms GPU
    forall rep. Body rep -> Stms rep
    bodyStms Body GPU
    body)
          [SubExpRes] -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {a} {c}.
    (Monad m, Bifunctor p, FreeIn a) =>
    a -> StateT (p Operands c) m ()
    collectFree (Body GPU -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult Body GPU
    body)
        collectStms :: Stms GPU -> State (Operands, Set VName) ()
    collectStms = (Stm GPU -> State (Operands, Set VName) ())
    -> Stms GPU -> State (Operands, Set VName) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ Stm GPU -> State (Operands, Set VName) ()
    collectStm

        collectStm :: Stm GPU -> State (Operands, Set VName) ()
    collectStm (Let Pat (LetDec GPU)
    pat StmAux (ExpDec GPU)
    _ Exp GPU
    ua)
          | BasicOp UpdateAcc {} <- Exp GPU
    ua,
            Pat [PatElem (LetDec GPU)
    pe] <- Pat (LetDec GPU)
    pat,
            Acc VName
    a ShapeBase SubExp
    _ [Type]
    _ NoUniqueness
    _ <- PatElem (LetDec GPU) -> Type
    forall t. Typed t => t -> Type
    typeOf PatElem (LetDec GPU)
    pe =
              -- Capture the tokens of accumulators used on host.
              VName -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {a} {a}.
    (Monad m, Bifunctor p, Ord a) =>
    a -> StateT (p a (Set a)) m ()
    captureAcc VName
    a State (Operands, Set VName) ()
    -> State (Operands, Set VName) () -> State (Operands, Set VName) ()
    forall a b.
    StateT (Operands, Set VName) Identity a
    -> StateT (Operands, Set VName) Identity b
    -> StateT (Operands, Set VName) Identity b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Exp GPU -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {rep} {c}.
    (Monad m, Bifunctor p) =>
    Exp rep -> StateT (p Operands c) m ()
    collectBasic Exp GPU
    ua
        collectStm Stm GPU
    stm = Exp GPU -> State (Operands, Set VName) ()
    collect (Stm GPU -> Exp GPU
    forall rep. Stm rep -> Exp rep
    stmExp Stm GPU
    stm)

        collectLForm :: LoopForm -> StateT (p Operands c) m ()
    collectLForm (ForLoop VName
    _ IntType
    _ SubExp
    b) = SubExp -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp SubExp
    b
        -- WhileLoop condition is declared as a loop parameter.
        collectLForm (WhileLoop VName
    _) = () -> StateT (p Operands c) m ()
    forall a. a -> StateT (p Operands c) m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

        -- The collective operands of an operator lambda body are only used on host
        -- if the associated accumulator is used in an UpdateAcc statement outside a
        -- kernel.
        collectWithAcc :: [WithAccInput GPU] -> Lambda GPU -> State (Operands, Set VName) ()
    collectWithAcc [WithAccInput GPU]
    inputs Lambda GPU
    f = do
          Body GPU -> State (Operands, Set VName) ()
    collectBody (Lambda GPU -> Body GPU
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPU
    f)
          Set VName
    used_accs <- ((Operands, Set VName) -> Set VName)
    -> StateT (Operands, Set VName) Identity (Set VName)
    forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
    gets (Operands, Set VName) -> Set VName
    forall a b. (a, b) -> b
    snd
          let accs :: [Type]
    accs = Id -> [Type] -> [Type]
    forall a. Id -> [a] -> [a]
    take ([WithAccInput GPU] -> Id
    forall a. [a] -> Id
    forall (t :: * -> *) a. Foldable t => t a -> Id
    length [WithAccInput GPU]
    inputs) (Lambda GPU -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda GPU
    f)
          let used :: [Bool]
    used = (Type -> Bool) -> [Type] -> [Bool]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(Acc VName
    a ShapeBase SubExp
    _ [Type]
    _ NoUniqueness
    _) -> VName -> Set VName -> Bool
    forall a. Ord a => a -> Set a -> Bool
    S.member VName
    a Set VName
    used_accs) [Type]
    accs
          ((Bool, WithAccInput GPU) -> State (Operands, Set VName) ())
    -> [(Bool, WithAccInput GPU)] -> State (Operands, Set VName) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Bool, WithAccInput GPU) -> State (Operands, Set VName) ()
    collectAcc ([Bool] -> [WithAccInput GPU] -> [(Bool, WithAccInput GPU)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Bool]
    used [WithAccInput GPU]
    inputs)

        collectAcc :: (Bool, WithAccInput GPU) -> State (Operands, Set VName) ()
    collectAcc (Bool
    _, (ShapeBase SubExp
    _, [VName]
    _, Maybe (Lambda GPU, [SubExp])
    Nothing)) = () -> State (Operands, Set VName) ()
    forall a. a -> StateT (Operands, Set VName) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
        collectAcc (Bool
    used, (ShapeBase SubExp
    _, [VName]
    _, Just (Lambda GPU
    op, [SubExp]
    nes))) = do
          (SubExp -> State (Operands, Set VName) ())
    -> [SubExp] -> State (Operands, Set VName) ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SubExp -> State (Operands, Set VName) ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp [SubExp]
    nes
          Bool
    -> State (Operands, Set VName) () -> State (Operands, Set VName) ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when Bool
    used (State (Operands, Set VName) () -> State (Operands, Set VName) ())
    -> State (Operands, Set VName) () -> State (Operands, Set VName) ()
    forall a b. (a -> b) -> a -> b
    $ Body GPU -> State (Operands, Set VName) ()
    collectBody (Lambda GPU -> Body GPU
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda GPU
    op)

        -- Does not collect named operands in
        --
        --   * types and shapes; size variables are assumed available to the host.
        --
        --   * use by a kernel body.
        --
        -- All other operands are conservatively collected even if they generally
        -- appear to be size variables or results computed by a SizeOp.
        collectHostOp :: HostOp op rep -> StateT (Operands, c) Identity ()
    collectHostOp (SegOp (SegMap SegLevel
    lvl SegSpace
    sp [Type]
    _ KernelBody rep
    _)) = do
          SegLevel -> StateT (Operands, c) Identity ()
    forall {c}. SegLevel -> StateT (Operands, c) Identity ()
    collectSegLevel SegLevel
    lvl
          SegSpace -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SegSpace -> StateT (p Operands c) m ()
    collectSegSpace SegSpace
    sp
        collectHostOp (SegOp (SegRed SegLevel
    lvl SegSpace
    sp [Type]
    _ KernelBody rep
    _ [SegBinOp rep]
    ops)) = do
          SegLevel -> StateT (Operands, c) Identity ()
    forall {c}. SegLevel -> StateT (Operands, c) Identity ()
    collectSegLevel SegLevel
    lvl
          SegSpace -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SegSpace -> StateT (p Operands c) m ()
    collectSegSpace SegSpace
    sp
          (SegBinOp rep -> StateT (Operands, c) Identity ())
    -> [SegBinOp rep] -> StateT (Operands, c) Identity ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SegBinOp rep -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {rep} {c}.
    (Monad m, Bifunctor p) =>
    SegBinOp rep -> StateT (p Operands c) m ()
    collectSegBinOp [SegBinOp rep]
    ops
        collectHostOp (SegOp (SegScan SegLevel
    lvl SegSpace
    sp [Type]
    _ KernelBody rep
    _ [SegBinOp rep]
    ops)) = do
          SegLevel -> StateT (Operands, c) Identity ()
    forall {c}. SegLevel -> StateT (Operands, c) Identity ()
    collectSegLevel SegLevel
    lvl
          SegSpace -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SegSpace -> StateT (p Operands c) m ()
    collectSegSpace SegSpace
    sp
          (SegBinOp rep -> StateT (Operands, c) Identity ())
    -> [SegBinOp rep] -> StateT (Operands, c) Identity ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SegBinOp rep -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {rep} {c}.
    (Monad m, Bifunctor p) =>
    SegBinOp rep -> StateT (p Operands c) m ()
    collectSegBinOp [SegBinOp rep]
    ops
        collectHostOp (SegOp (SegHist SegLevel
    lvl SegSpace
    sp [Type]
    _ KernelBody rep
    _ [HistOp rep]
    ops)) = do
          SegLevel -> StateT (Operands, c) Identity ()
    forall {c}. SegLevel -> StateT (Operands, c) Identity ()
    collectSegLevel SegLevel
    lvl
          SegSpace -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SegSpace -> StateT (p Operands c) m ()
    collectSegSpace SegSpace
    sp
          (HistOp rep -> StateT (Operands, c) Identity ())
    -> [HistOp rep] -> StateT (Operands, c) Identity ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ HistOp rep -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {rep} {c}.
    (Monad m, Bifunctor p) =>
    HistOp rep -> StateT (p Operands c) m ()
    collectHistOp [HistOp rep]
    ops
        collectHostOp (SizeOp SizeOp
    op) = SizeOp -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {a} {c}.
    (Monad m, Bifunctor p, FreeIn a) =>
    a -> StateT (p Operands c) m ()
    collectFree SizeOp
    op
        collectHostOp (OtherOp op rep
    op) = op rep -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {a} {c}.
    (Monad m, Bifunctor p, FreeIn a) =>
    a -> StateT (p Operands c) m ()
    collectFree op rep
    op
        collectHostOp GPUBody {} = () -> StateT (Operands, c) Identity ()
    forall a. a -> StateT (Operands, c) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

        collectSegLevel :: SegLevel -> StateT (Operands, c) Identity ()
    collectSegLevel = (VName -> StateT (Operands, c) Identity ())
    -> [VName] -> StateT (Operands, c) Identity ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ VName -> StateT (Operands, c) Identity ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    VName -> StateT (p Operands c) m ()
    captureName ([VName] -> StateT (Operands, c) Identity ())
    -> (SegLevel -> [VName])
    -> SegLevel
    -> StateT (Operands, c) Identity ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Names -> [VName]
    namesToList (Names -> [VName]) -> (SegLevel -> Names) -> SegLevel -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegLevel -> Names
    forall a. FreeIn a => a -> Names
    freeIn

        collectSegSpace :: SegSpace -> StateT (p Operands c) m ()
    collectSegSpace SegSpace
    space =
          (SubExp -> StateT (p Operands c) m ())
    -> [SubExp] -> StateT (p Operands c) m ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SubExp -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp (SegSpace -> [SubExp]
    segSpaceDims SegSpace
    space)

        collectSegBinOp :: SegBinOp rep -> StateT (p Operands c) m ()
    collectSegBinOp (SegBinOp Commutativity
    _ Lambda rep
    _ [SubExp]
    nes ShapeBase SubExp
    _) =
          (SubExp -> StateT (p Operands c) m ())
    -> [SubExp] -> StateT (p Operands c) m ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SubExp -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp [SubExp]
    nes

        collectHistOp :: HistOp rep -> StateT (p Operands c) m ()
    collectHistOp (HistOp ShapeBase SubExp
    _ SubExp
    rf [VName]
    _ [SubExp]
    nes ShapeBase SubExp
    _ Lambda rep
    _) = do
          SubExp -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp SubExp
    rf
          (SubExp -> StateT (p Operands c) m ())
    -> [SubExp] -> StateT (p Operands c) m ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ SubExp -> StateT (p Operands c) m ()
    forall {m :: * -> *} {p :: * -> * -> *} {c}.
    (Monad m, Bifunctor p) =>
    SubExp -> StateT (p Operands c) m ()
    collectSubExp [SubExp]
    nes

    --------------------------------------------------------------------------------
    --                        GRAPH BUILDING - PRIMITIVES                         --
    --------------------------------------------------------------------------------

    -- | Creates a vertex for the given binding, provided that the set of operands
    -- is not empty.
    createNode :: Binding -> Operands -> Grapher ()
    createNode :: Binding -> Operands -> Grapher ()
    createNode Binding
    b Operands
    ops =
      Bool -> Grapher () -> Grapher ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (Operands -> Bool
    IS.null Operands
    ops) (Binding -> Grapher ()
    addVertex Binding
    b Grapher () -> Grapher () -> Grapher ()
    forall a b.
    StateT State (Reader Env) a
    -> StateT State (Reader Env) b -> StateT State (Reader Env) b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Edges -> Operands -> Grapher ()
    addEdges (Id -> Edges
    MG.oneEdge (Id -> Edges) -> Id -> Edges
    forall a b. (a -> b) -> a -> b
    $ Binding -> Id
    forall a b. (a, b) -> a
    fst Binding
    b) Operands
    ops)

    -- | Adds a vertex to the graph for the given binding.
    addVertex :: Binding -> Grapher ()
    addVertex :: Binding -> Grapher ()
    addVertex (Id
    i, Type
    t) = do
      Meta
    meta <- Grapher Meta
    getMeta
      let v :: Vertex Meta
    v = Id -> Meta -> Vertex Meta
    forall m. Id -> m -> Vertex m
    MG.vertex Id
    i Meta
    meta
      Bool -> Grapher () -> Grapher ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when (Type -> Bool
    forall t. Typed t => t -> Bool
    isScalar Type
    t) (Grapher () -> Grapher ()) -> Grapher () -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ (Operands -> Operands) -> Grapher ()
    modifyGraphedScalars (Id -> Operands -> Operands
    IS.insert Id
    i)
      Bool -> Grapher () -> Grapher ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    when (Type -> Bool
    forall t. Typed t => t -> Bool
    isArray Type
    t) (Grapher () -> Grapher ()) -> Grapher () -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ Id -> Id -> Grapher ()
    recordCopyableMemory Id
    i (Meta -> Id
    metaBodyDepth Meta
    meta)
      (Graph -> Graph) -> Grapher ()
    modifyGraph (Vertex Meta -> Graph -> Graph
    forall m. Vertex m -> Graph m -> Graph m
    MG.insert Vertex Meta
    v)

    -- | Adds a source connected vertex to the graph for the given binding.
    addSource :: Binding -> Grapher ()
    addSource :: Binding -> Grapher ()
    addSource Binding
    b = do
      Binding -> Grapher ()
    addVertex Binding
    b
      (Sources -> Sources) -> Grapher ()
    modifySources ((Sources -> Sources) -> Grapher ())
    -> (Sources -> Sources) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ ([Id] -> [Id]) -> Sources -> Sources
    forall b c a. (b -> c) -> (a, b) -> (a, c)
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second (Binding -> Id
    forall a b. (a, b) -> a
    fst Binding
    b :)

    -- | Adds the given edges to each vertex identified by the 'IdSet'. It is
    -- assumed that all vertices reside within the body that currently is being
    -- graphed.
    addEdges :: Edges -> IdSet -> Grapher ()
    addEdges :: Edges -> Operands -> Grapher ()
    addEdges Edges
    ToSink Operands
    is = do
      (Graph -> Graph) -> Grapher ()
    modifyGraph ((Graph -> Graph) -> Grapher ()) -> (Graph -> Graph) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \Graph
    g -> (Graph -> Id -> Graph) -> Graph -> Operands -> Graph
    forall a. (a -> Id -> a) -> a -> Operands -> a
    IS.foldl' ((Id -> Graph -> Graph) -> Graph -> Id -> Graph
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Id -> Graph -> Graph
    forall m. Id -> Graph m -> Graph m
    MG.connectToSink) Graph
    g Operands
    is
      (Operands -> Operands) -> Grapher ()
    modifyGraphedScalars (Operands -> Operands -> Operands
    `IS.difference` Operands
    is)
    addEdges Edges
    es Operands
    is = do
      (Graph -> Graph) -> Grapher ()
    modifyGraph ((Graph -> Graph) -> Grapher ()) -> (Graph -> Graph) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \Graph
    g -> (Graph -> Id -> Graph) -> Graph -> Operands -> Graph
    forall a. (a -> Id -> a) -> a -> Operands -> a
    IS.foldl' ((Id -> Graph -> Graph) -> Graph -> Id -> Graph
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip ((Id -> Graph -> Graph) -> Graph -> Id -> Graph)
    -> (Id -> Graph -> Graph) -> Graph -> Id -> Graph
    forall a b. (a -> b) -> a -> b
    $ Edges -> Id -> Graph -> Graph
    forall m. Edges -> Id -> Graph m -> Graph m
    MG.addEdges Edges
    es) Graph
    g Operands
    is
      Operands -> Grapher ()
    tellOperands Operands
    is

    -- | Ensure that a variable (which is in scope) will be made available on host
    -- before its first use.
    requiredOnHost :: Id -> Grapher ()
    requiredOnHost :: Id -> Grapher ()
    requiredOnHost Id
    i = do
      Maybe (Vertex Meta)
    mv <- Id -> Graph -> Maybe (Vertex Meta)
    forall m. Id -> Graph m -> Maybe (Vertex m)
    MG.lookup Id
    i (Graph -> Maybe (Vertex Meta))
    -> Grapher Graph -> StateT State (Reader Env) (Maybe (Vertex Meta))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Grapher Graph
    getGraph
      case Maybe (Vertex Meta)
    mv of
        Maybe (Vertex Meta)
    Nothing -> () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()
        Just Vertex Meta
    v -> do
          Id -> Grapher ()
    connectToSink Id
    i
          Id -> Grapher ()
    tellHostOnlyParent (Meta -> Id
    metaBodyDepth (Meta -> Id) -> Meta -> Id
    forall a b. (a -> b) -> a -> b
    $ Vertex Meta -> Meta
    forall m. Vertex m -> m
    vertexMeta Vertex Meta
    v)

    -- | Connects the vertex of the given id to a sink.
    connectToSink :: Id -> Grapher ()
    connectToSink :: Id -> Grapher ()
    connectToSink Id
    i = do
      (Graph -> Graph) -> Grapher ()
    modifyGraph (Id -> Graph -> Graph
    forall m. Id -> Graph m -> Graph m
    MG.connectToSink Id
    i)
      (Operands -> Operands) -> Grapher ()
    modifyGraphedScalars (Id -> Operands -> Operands
    IS.delete Id
    i)

    -- | Like 'connectToSink' but vertex is given by a t'SubExp'. This is a no-op if
    -- the t'SubExp' is a constant.
    connectSubExpToSink :: SubExp -> Grapher ()
    connectSubExpToSink :: SubExp -> Grapher ()
    connectSubExpToSink (Var VName
    n) = Id -> Grapher ()
    connectToSink (VName -> Id
    nameToId VName
    n)
    connectSubExpToSink SubExp
    _ = () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

    -- | Routes all possible routes within the subgraph identified by this id.
    -- Returns the ids of the source connected vertices that were attempted routed.
    --
    -- Assumption: The subgraph with the given id has just been created and no path
    -- exists from it to an external sink.
    routeSubgraph :: Id -> Grapher [Id]
    routeSubgraph :: Id -> Grapher [Id]
    routeSubgraph Id
    si = do
      State
    st <- StateT State (Reader Env) State
    forall (m :: * -> *) s. Monad m => StateT s m s
    get
      let g :: Graph
    g = State -> Graph
    stateGraph State
    st
      let ([Id]
    routed, [Id]
    unrouted) = State -> Sources
    stateSources State
    st
      let ([Id]
    gsrcs, [Id]
    unrouted') = (Id -> Bool) -> [Id] -> Sources
    forall a. (a -> Bool) -> [a] -> ([a], [a])
    span (Id -> Graph -> Id -> Bool
    inSubGraph Id
    si Graph
    g) [Id]
    unrouted
      let ([Id]
    sinks, Graph
    g') = [Id] -> Graph -> ([Id], Graph)
    forall m. [Id] -> Graph m -> ([Id], Graph m)
    MG.routeMany [Id]
    gsrcs Graph
    g
      State -> Grapher ()
    forall (m :: * -> *) s. Monad m => s -> StateT s m ()
    put (State -> Grapher ()) -> State -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $
        State
    st
          { stateGraph = g',
            stateSources = (gsrcs ++ routed, unrouted'),
            stateSinks = sinks ++ stateSinks st
          }
      [Id] -> Grapher [Id]
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [Id]
    gsrcs

    -- | @inSubGraph si g i@ returns whether @g@ contains a vertex with id @i@ that
    -- is declared within the subgraph with id @si@.
    inSubGraph :: Id -> Graph -> Id -> Bool
    inSubGraph :: Id -> Graph -> Id -> Bool
    inSubGraph Id
    si Graph
    g Id
    i
      | Just Vertex Meta
    v <- Id -> Graph -> Maybe (Vertex Meta)
    forall m. Id -> Graph m -> Maybe (Vertex m)
    MG.lookup Id
    i Graph
    g,
        Just Id
    mgi <- Meta -> Maybe Id
    metaGraphId (Vertex Meta -> Meta
    forall m. Vertex m -> m
    vertexMeta Vertex Meta
    v) =
          Id
    si Id -> Id -> Bool
    forall a. Eq a => a -> a -> Bool
    == Id
    mgi
    inSubGraph Id
    _ Graph
    _ Id
    _ = Bool
    False

    -- | @b `reuses` n@ records that @b@ binds an array backed by the same memory
    -- as @n@. If @b@ is not array typed or the backing memory is not copyable then
    -- this does nothing.
    reuses :: Binding -> VName -> Grapher ()
    reuses :: Binding -> VName -> Grapher ()
    reuses (Id
    i, Type
    t) VName
    n
      | Type -> Bool
    forall t. Typed t => t -> Bool
    isArray Type
    t =
          do
            Maybe Id
    body_depth <- VName -> Grapher (Maybe Id)
    outermostCopyableArray VName
    n
            Maybe Id -> (Id -> Grapher ()) -> Grapher ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ Maybe Id
    body_depth (Id -> Id -> Grapher ()
    recordCopyableMemory Id
    i)
      | Bool
    otherwise =
          () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

    reusesSubExp :: Binding -> SubExp -> Grapher ()
    reusesSubExp :: Binding -> SubExp -> Grapher ()
    reusesSubExp Binding
    b (Var VName
    n) = Binding
    b Binding -> VName -> Grapher ()
    `reuses` VName
    n
    reusesSubExp Binding
    _ SubExp
    _ = () -> Grapher ()
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()

    -- @reusesReturn bs res@ records each array binding in @bs@ as reusing copyable
    -- memory if the corresponding return value in @res@ is backed by copyable
    -- memory.
    --
    -- If every array binding is registered as being backed by copyable memory then
    -- the function returns @True@, otherwise it returns @False@.
    reusesReturn :: [Binding] -> [SubExp] -> Grapher Bool
    reusesReturn :: [Binding] -> [SubExp] -> Grapher Bool
    reusesReturn [Binding]
    bs [SubExp]
    res = do
      Id
    body_depth <- Meta -> Id
    metaBodyDepth (Meta -> Id) -> Grapher Meta -> StateT State (Reader Env) Id
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Grapher Meta
    getMeta
      (Bool -> (Binding, SubExp) -> Grapher Bool)
    -> Bool -> [(Binding, SubExp)] -> Grapher Bool
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM (Id -> Bool -> (Binding, SubExp) -> Grapher Bool
    reuse Id
    body_depth) Bool
    True ([Binding] -> [SubExp] -> [(Binding, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Binding]
    bs [SubExp]
    res)
      where
        reuse :: Int -> Bool -> (Binding, SubExp) -> Grapher Bool
        reuse :: Id -> Bool -> (Binding, SubExp) -> Grapher Bool
    reuse Id
    body_depth Bool
    onlyCopyable (Binding
    b, SubExp
    se)
          | (SubExp -> Bool) -> [SubExp] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (SubExp -> SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1) (Type -> [SubExp]
    forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
    arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Binding -> Type
    forall a b. (a, b) -> b
    snd Binding
    b) =
              -- Single element arrays are immediately recognizable as copyable so
              -- don't bother recording those. Note that this case also matches
              -- primitive return values.
              Bool -> Grapher Bool
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Bool
    onlyCopyable
          | (Id
    i, Type
    t) <- Binding
    b,
            Type -> Bool
    forall t. Typed t => t -> Bool
    isArray Type
    t,
            Var VName
    n <- SubExp
    se =
              do
                Maybe Id
    res_body_depth <- VName -> Grapher (Maybe Id)
    outermostCopyableArray VName
    n
                case Maybe Id
    res_body_depth of
                  Just Id
    inner -> do
                    Id -> Id -> Grapher ()
    recordCopyableMemory Id
    i (Id -> Id -> Id
    forall a. Ord a => a -> a -> a
    min Id
    body_depth Id
    inner)
                    let returns_free_var :: Bool
    returns_free_var = Id
    inner Id -> Id -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Id
    body_depth
                    Bool -> Grapher Bool
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Bool
    onlyCopyable Bool -> Bool -> Bool
    && Bool -> Bool
    not Bool
    returns_free_var)
                  Maybe Id
    _ ->
                    Bool -> Grapher Bool
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Bool
    False
          | Bool
    otherwise =
              Bool -> Grapher Bool
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Bool
    onlyCopyable

    -- @reusesBranches bs seses@ records each array binding in @bs@ as
    -- reusing copyable memory if each corresponding return value in the
    -- lists in @ses@ are backed by copyable memory.  Each list is the
    -- result of a branch body (i.e. for 'if' the list has two elements).
    --
    -- If every array binding is registered as being backed by copyable
    -- memory then the function returns @True@, otherwise it returns
    -- @False@.
    reusesBranches :: [Binding] -> [[SubExp]] -> Grapher Bool
    reusesBranches :: [Binding] -> [[SubExp]] -> Grapher Bool
    reusesBranches [Binding]
    bs [[SubExp]]
    seses = do
      Id
    body_depth <- Meta -> Id
    metaBodyDepth (Meta -> Id) -> Grapher Meta -> StateT State (Reader Env) Id
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Grapher Meta
    getMeta
      (Bool -> (Binding, [SubExp]) -> Grapher Bool)
    -> Bool -> [(Binding, [SubExp])] -> Grapher Bool
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM (Id -> Bool -> (Binding, [SubExp]) -> Grapher Bool
    reuse Id
    body_depth) Bool
    True ([(Binding, [SubExp])] -> Grapher Bool)
    -> [(Binding, [SubExp])] -> Grapher Bool
    forall a b. (a -> b) -> a -> b
    $ [Binding] -> [[SubExp]] -> [(Binding, [SubExp])]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Binding]
    bs ([[SubExp]] -> [(Binding, [SubExp])])
    -> [[SubExp]] -> [(Binding, [SubExp])]
    forall a b. (a -> b) -> a -> b
    $ [[SubExp]] -> [[SubExp]]
    forall a. [[a]] -> [[a]]
    L.transpose [[SubExp]]
    seses
      where
        reuse :: Int -> Bool -> (Binding, [SubExp]) -> Grapher Bool
        reuse :: Id -> Bool -> (Binding, [SubExp]) -> Grapher Bool
    reuse Id
    body_depth Bool
    onlyCopyable (Binding
    b, [SubExp]
    ses)
          | (SubExp -> Bool) -> [SubExp] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (SubExp -> SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1) (Type -> [SubExp]
    forall u. TypeBase (ShapeBase SubExp) u -> [SubExp]
    arrayDims (Type -> [SubExp]) -> Type -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Binding -> Type
    forall a b. (a, b) -> b
    snd Binding
    b) =
              -- Single element arrays are immediately recognizable as copyable so
              -- don't bother recording those. Note that this case also matches
              -- primitive return values.
              Bool -> Grapher Bool
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Bool
    onlyCopyable
          | (Id
    i, Type
    t) <- Binding
    b,
            Type -> Bool
    forall t. Typed t => t -> Bool
    isArray Type
    t,
            Just [VName]
    ns <- (SubExp -> Maybe VName) -> [SubExp] -> Maybe [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 SubExp -> Maybe VName
    subExpVar [SubExp]
    ses = do
              [Maybe Id]
    body_depths <- (VName -> Grapher (Maybe Id))
    -> [VName] -> StateT State (Reader Env) [Maybe Id]
    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 -> Grapher (Maybe Id)
    outermostCopyableArray [VName]
    ns
              case [Maybe Id] -> Maybe [Id]
    forall (t :: * -> *) (m :: * -> *) a.
    (Traversable t, Monad m) =>
    t (m a) -> m (t a)
    forall (m :: * -> *) a. Monad m => [m a] -> m [a]
    sequence [Maybe Id]
    body_depths of
                Just [Id]
    bds -> do
                  let inner :: Id
    inner = [Id] -> Id
    forall a. Ord a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
    minimum [Id]
    bds
                  Id -> Id -> Grapher ()
    recordCopyableMemory Id
    i (Id -> Id -> Id
    forall a. Ord a => a -> a -> a
    min Id
    body_depth Id
    inner)
                  let returns_free_var :: Bool
    returns_free_var = Id
    inner Id -> Id -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Id
    body_depth
                  Bool -> Grapher Bool
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Bool
    onlyCopyable Bool -> Bool -> Bool
    && Bool -> Bool
    not Bool
    returns_free_var)
                Maybe [Id]
    _ ->
                  Bool -> Grapher Bool
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Bool
    False
          | Bool
    otherwise =
              Bool -> Grapher Bool
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Bool
    onlyCopyable

    --------------------------------------------------------------------------------
    --                           GRAPH BUILDING - TYPES                           --
    --------------------------------------------------------------------------------

    type Grapher = StateT State (R.Reader Env)

    data Env = Env
      { -- | See 'HostOnlyFuns'.
        Env -> HostOnlyFuns
    envHostOnlyFuns :: HostOnlyFuns,
        -- | Metadata for the current body being graphed.
        Env -> Meta
    envMeta :: Meta
      }

    -- | A measurement of how many bodies something is nested within.
    type BodyDepth = Int

    -- | Metadata on the environment that a variable is declared within.
    data Meta = Meta
      { -- | How many if statement branch bodies the variable binding is nested
        -- within. If a route passes through the edge u->v and the fork depth
        --
        --   1) increases from u to v, then u is within a conditional branch.
        --
        --   2) decreases from u to v, then v binds the result of two or more
        --      branches.
        --
        -- After the graph has been built and routed, this can be used to delay
        -- reads into deeper branches to reduce their likelihood of manifesting.
        Meta -> Id
    metaForkDepth :: Int,
        -- | How many bodies the variable is nested within.
        Meta -> Id
    metaBodyDepth :: BodyDepth,
        -- | An id for the subgraph within which the variable exists, defined at
        -- the body level. A read may only be delayed to a point within its own
        -- subgraph.
        Meta -> Maybe Id
    metaGraphId :: Maybe Id
      }

    -- | Ids for all variables used as an operand.
    type Operands = IdSet

    -- | Statistics on the statements within a body and their dependencies.
    data BodyStats = BodyStats
      { -- | Whether the body contained any host-only statements.
        BodyStats -> Bool
    bodyHostOnly :: Bool,
        -- | Whether the body contained any GPUBody kernels.
        BodyStats -> Bool
    bodyHasGPUBody :: Bool,
        -- | Whether the body performed any reads.
        BodyStats -> Bool
    bodyReads :: Bool,
        -- | All scalar variables represented in the graph that have been used
        -- as return values of the body or as operands within it, including those
        -- that are defined within the body itself. Variables with vertices
        -- connected to sinks may be excluded.
        BodyStats -> Operands
    bodyOperands :: Operands,
        -- | Depth of parent bodies with variables that are required on host. Since
        -- the variables are required on host, the parent statements of these bodies
        -- cannot be moved to device as a whole. They are host-only.
        BodyStats -> Operands
    bodyHostOnlyParents :: IS.IntSet
      }

    instance Semigroup BodyStats where
      (BodyStats Bool
    ho1 Bool
    gb1 Bool
    r1 Operands
    o1 Operands
    hop1) <> :: BodyStats -> BodyStats -> BodyStats
    <> (BodyStats Bool
    ho2 Bool
    gb2 Bool
    r2 Operands
    o2 Operands
    hop2) =
        BodyStats
          { bodyHostOnly :: Bool
    bodyHostOnly = Bool
    ho1 Bool -> Bool -> Bool
    || Bool
    ho2,
            bodyHasGPUBody :: Bool
    bodyHasGPUBody = Bool
    gb1 Bool -> Bool -> Bool
    || Bool
    gb2,
            bodyReads :: Bool
    bodyReads = Bool
    r1 Bool -> Bool -> Bool
    || Bool
    r2,
            bodyOperands :: Operands
    bodyOperands = Operands -> Operands -> Operands
    IS.union Operands
    o1 Operands
    o2,
            bodyHostOnlyParents :: Operands
    bodyHostOnlyParents = Operands -> Operands -> Operands
    IS.union Operands
    hop1 Operands
    hop2
          }

    instance Monoid BodyStats where
      mempty :: BodyStats
    mempty =
        BodyStats
          { bodyHostOnly :: Bool
    bodyHostOnly = Bool
    False,
            bodyHasGPUBody :: Bool
    bodyHasGPUBody = Bool
    False,
            bodyReads :: Bool
    bodyReads = Bool
    False,
            bodyOperands :: Operands
    bodyOperands = Operands
    IS.empty,
            bodyHostOnlyParents :: Operands
    bodyHostOnlyParents = Operands
    IS.empty
          }

    type Graph = MG.Graph Meta

    -- | All vertices connected from a source, partitioned into those that have
    -- been attempted routed and those which have not.
    type Sources = ([Id], [Id])

    -- | All terminal vertices of routes.
    type Sinks = [Id]

    -- | A captured statement for which graphing has been delayed.
    type Delayed = (Binding, Exp GPU)

    -- | The vertex handle for a variable and its type.
    type Binding = (Id, Type)

    -- | Array variables backed by memory segments that may be copied, mapped to the
    -- outermost known body depths that declares arrays backed by a superset of
    -- those segments.
    type CopyableMemoryMap = IM.IntMap BodyDepth

    data State = State
      { -- | The graph being built.
        State -> Graph
    stateGraph :: Graph,
        -- | All known scalars that have been graphed.
        State -> Operands
    stateGraphedScalars :: IdSet,
        -- | All variables that directly bind scalars read from device memory.
        State -> Sources
    stateSources :: Sources,
        -- | Graphed scalars that are used as operands by statements that cannot be
        -- migrated. A read cannot be delayed beyond these, so if the statements
        -- that bind these variables are moved to device, the variables must be read
        -- from device memory.
        State -> [Id]
    stateSinks :: Sinks,
        -- | Observed 'UpdateAcc' host statements to be graphed later.
        State -> IntMap [Delayed]
    stateUpdateAccs :: IM.IntMap [Delayed],
        -- | A map of encountered arrays that are backed by copyable memory.
        -- Trivial instances such as single element arrays are excluded.
        State -> CopyableMemoryMap
    stateCopyableMemory :: CopyableMemoryMap,
        -- | Information about the current body being graphed.
        State -> BodyStats
    stateStats :: BodyStats
      }

    --------------------------------------------------------------------------------
    --                             GRAPHER OPERATIONS                             --
    --------------------------------------------------------------------------------

    execGrapher :: HostOnlyFuns -> Grapher a -> (Graph, Sources, Sinks)
    execGrapher :: forall a. HostOnlyFuns -> Grapher a -> (Graph, Sources, [Id])
    execGrapher HostOnlyFuns
    hof Grapher a
    m =
      let s :: State
    s = Reader Env State -> Env -> State
    forall r a. Reader r a -> r -> a
    R.runReader (Grapher a -> State -> Reader Env State
    forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
    execStateT Grapher a
    m State
    st) Env
    env
       in (State -> Graph
    stateGraph State
    s, State -> Sources
    stateSources State
    s, State -> [Id]
    stateSinks State
    s)
      where
        env :: Env
    env =
          Env
            { envHostOnlyFuns :: HostOnlyFuns
    envHostOnlyFuns = HostOnlyFuns
    hof,
              envMeta :: Meta
    envMeta =
                Meta
                  { metaForkDepth :: Id
    metaForkDepth = Id
    0,
                    metaBodyDepth :: Id
    metaBodyDepth = Id
    0,
                    metaGraphId :: Maybe Id
    metaGraphId = Maybe Id
    forall a. Maybe a
    Nothing
                  }
            }
        st :: State
    st =
          State
            { stateGraph :: Graph
    stateGraph = Graph
    forall m. Graph m
    MG.empty,
              stateGraphedScalars :: Operands
    stateGraphedScalars = Operands
    IS.empty,
              stateSources :: Sources
    stateSources = ([], []),
              stateSinks :: [Id]
    stateSinks = [],
              stateUpdateAccs :: IntMap [Delayed]
    stateUpdateAccs = IntMap [Delayed]
    forall a. IntMap a
    IM.empty,
              stateCopyableMemory :: CopyableMemoryMap
    stateCopyableMemory = CopyableMemoryMap
    forall a. IntMap a
    IM.empty,
              stateStats :: BodyStats
    stateStats = BodyStats
    forall a. Monoid a => a
    mempty
            }

    -- | Execute a computation in a modified environment.
    local :: (Env -> Env) -> Grapher a -> Grapher a
    local :: forall a. (Env -> Env) -> Grapher a -> Grapher a
    local Env -> Env
    f = (Reader Env (a, State) -> Reader Env (a, State))
    -> StateT State (Reader Env) a -> StateT State (Reader Env) a
    forall (m :: * -> *) a s (n :: * -> *) b.
    (m (a, s) -> n (b, s)) -> StateT s m a -> StateT s n b
    mapStateT ((Env -> Env) -> Reader Env (a, State) -> Reader Env (a, State)
    forall r (m :: * -> *) a.
    (r -> r) -> ReaderT r m a -> ReaderT r m a
    R.local Env -> Env
    f)

    -- | Fetch the value of the environment.
    ask :: Grapher Env
    ask :: Grapher Env
    ask = Reader Env Env -> Grapher Env
    forall (m :: * -> *) a. Monad m => m a -> StateT State m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift Reader Env Env
    forall (m :: * -> *) r. Monad m => ReaderT r m r
    R.ask

    -- | Retrieve a function of the current environment.
    asks :: (Env -> a) -> Grapher a
    asks :: forall a. (Env -> a) -> Grapher a
    asks = Reader Env a -> StateT State (Reader Env) a
    forall (m :: * -> *) a. Monad m => m a -> StateT State m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift (Reader Env a -> StateT State (Reader Env) a)
    -> ((Env -> a) -> Reader Env a)
    -> (Env -> a)
    -> StateT State (Reader Env) a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Env -> a) -> Reader Env a
    forall (m :: * -> *) r a. Monad m => (r -> a) -> ReaderT r m a
    R.asks

    -- | Register that the body contains a host-only statement. This means its
    -- parent statement and any parent bodies themselves are host-only. A host-only
    -- statement should not be migrated, either because it cannot run on device or
    -- because it would be inefficient to do so.
    tellHostOnly :: Grapher ()
    tellHostOnly :: Grapher ()
    tellHostOnly =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateStats = (stateStats st) {bodyHostOnly = True}}

    -- | Register that the body contains a GPUBody kernel.
    tellGPUBody :: Grapher ()
    tellGPUBody :: Grapher ()
    tellGPUBody =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateStats = (stateStats st) {bodyHasGPUBody = True}}

    -- | Register that the current body contains a statement that reads device
    -- memory.
    tellRead :: Grapher ()
    tellRead :: Grapher ()
    tellRead =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateStats = (stateStats st) {bodyReads = True}}

    -- | Register that these variables are used as operands within the current body.
    tellOperands :: IdSet -> Grapher ()
    tellOperands :: Operands -> Grapher ()
    tellOperands Operands
    is =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st ->
        let stats :: BodyStats
    stats = State -> BodyStats
    stateStats State
    st
            operands :: Operands
    operands = BodyStats -> Operands
    bodyOperands BodyStats
    stats
         in State
    st {stateStats = stats {bodyOperands = operands <> is}}

    -- | Register that the current statement with a body at the given body depth is
    -- host-only.
    tellHostOnlyParent :: BodyDepth -> Grapher ()
    tellHostOnlyParent :: Id -> Grapher ()
    tellHostOnlyParent Id
    body_depth =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st ->
        let stats :: BodyStats
    stats = State -> BodyStats
    stateStats State
    st
            parents :: Operands
    parents = BodyStats -> Operands
    bodyHostOnlyParents BodyStats
    stats
            parents' :: Operands
    parents' = Id -> Operands -> Operands
    IS.insert Id
    body_depth Operands
    parents
         in State
    st {stateStats = stats {bodyHostOnlyParents = parents'}}

    -- | Get the graph under construction.
    getGraph :: Grapher Graph
    getGraph :: Grapher Graph
    getGraph = (State -> Graph) -> Grapher Graph
    forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
    gets State -> Graph
    stateGraph

    -- | All scalar variables with a vertex representation in the graph.
    getGraphedScalars :: Grapher IdSet
    getGraphedScalars :: Grapher Operands
    getGraphedScalars = (State -> Operands) -> Grapher Operands
    forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
    gets State -> Operands
    stateGraphedScalars

    -- | Every known array that is backed by a memory segment that may be copied,
    -- mapped to the outermost known body depth where an array is backed by a
    -- superset of that segment.
    --
    -- A body where all returned arrays are backed by such memory and are written by
    -- its own statements will retain its asymptotic cost if migrated as a whole.
    getCopyableMemory :: Grapher CopyableMemoryMap
    getCopyableMemory :: Grapher CopyableMemoryMap
    getCopyableMemory = (State -> CopyableMemoryMap) -> Grapher CopyableMemoryMap
    forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
    gets State -> CopyableMemoryMap
    stateCopyableMemory

    -- | The outermost known body depth for an array backed by the same copyable
    -- memory as the array with this name.
    outermostCopyableArray :: VName -> Grapher (Maybe BodyDepth)
    outermostCopyableArray :: VName -> Grapher (Maybe Id)
    outermostCopyableArray VName
    n = Id -> CopyableMemoryMap -> Maybe Id
    forall a. Id -> IntMap a -> Maybe a
    IM.lookup (VName -> Id
    nameToId VName
    n) (CopyableMemoryMap -> Maybe Id)
    -> Grapher CopyableMemoryMap -> Grapher (Maybe Id)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Grapher CopyableMemoryMap
    getCopyableMemory

    -- | Reduces the variables to just the 'Id's of those that are scalars and which
    -- have a vertex representation in the graph, excluding those that have been
    -- connected to sinks.
    onlyGraphedScalars :: (Foldable t) => t VName -> Grapher IdSet
    onlyGraphedScalars :: forall (t :: * -> *). Foldable t => t VName -> Grapher Operands
    onlyGraphedScalars t VName
    vs = do
      let is :: Operands
    is = (Operands -> VName -> Operands) -> Operands -> t VName -> Operands
    forall b a. (b -> a -> b) -> b -> t a -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' (\Operands
    s VName
    n -> Id -> Operands -> Operands
    IS.insert (VName -> Id
    nameToId VName
    n) Operands
    s) Operands
    IS.empty t VName
    vs
      Operands -> Operands -> Operands
    IS.intersection Operands
    is (Operands -> Operands) -> Grapher Operands -> Grapher Operands
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Grapher Operands
    getGraphedScalars

    -- | Like 'onlyGraphedScalars' but for a single 'VName'.
    onlyGraphedScalar :: VName -> Grapher IdSet
    onlyGraphedScalar :: VName -> Grapher Operands
    onlyGraphedScalar VName
    n = do
      let i :: Id
    i = VName -> Id
    nameToId VName
    n
      Operands
    gss <- Grapher Operands
    getGraphedScalars
      if Id -> Operands -> Bool
    IS.member Id
    i Operands
    gss
        then Operands -> Grapher Operands
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Id -> Operands
    IS.singleton Id
    i)
        else Operands -> Grapher Operands
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Operands
    IS.empty

    -- | Like 'onlyGraphedScalars' but for a single t'SubExp'.
    onlyGraphedScalarSubExp :: SubExp -> Grapher IdSet
    onlyGraphedScalarSubExp :: SubExp -> Grapher Operands
    onlyGraphedScalarSubExp (Constant PrimValue
    _) = Operands -> Grapher Operands
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Operands
    IS.empty
    onlyGraphedScalarSubExp (Var VName
    n) = VName -> Grapher Operands
    onlyGraphedScalar VName
    n

    -- | Update the graph under construction.
    modifyGraph :: (Graph -> Graph) -> Grapher ()
    modifyGraph :: (Graph -> Graph) -> Grapher ()
    modifyGraph Graph -> Graph
    f =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateGraph = f (stateGraph st)}

    -- | Update the contents of the graphed scalar set.
    modifyGraphedScalars :: (IdSet -> IdSet) -> Grapher ()
    modifyGraphedScalars :: (Operands -> Operands) -> Grapher ()
    modifyGraphedScalars Operands -> Operands
    f =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateGraphedScalars = f (stateGraphedScalars st)}

    -- | Update the contents of the copyable memory map.
    modifyCopyableMemory :: (CopyableMemoryMap -> CopyableMemoryMap) -> Grapher ()
    modifyCopyableMemory :: (CopyableMemoryMap -> CopyableMemoryMap) -> Grapher ()
    modifyCopyableMemory CopyableMemoryMap -> CopyableMemoryMap
    f =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateCopyableMemory = f (stateCopyableMemory st)}

    -- | Update the set of source connected vertices.
    modifySources :: (Sources -> Sources) -> Grapher ()
    modifySources :: (Sources -> Sources) -> Grapher ()
    modifySources Sources -> Sources
    f =
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateSources = f (stateSources st)}

    -- | Record that this variable binds an array that is backed by copyable
    -- memory shared by an array at this outermost body depth.
    recordCopyableMemory :: Id -> BodyDepth -> Grapher ()
    recordCopyableMemory :: Id -> Id -> Grapher ()
    recordCopyableMemory Id
    i Id
    bd =
      (CopyableMemoryMap -> CopyableMemoryMap) -> Grapher ()
    modifyCopyableMemory (Id -> Id -> CopyableMemoryMap -> CopyableMemoryMap
    forall a. Id -> a -> IntMap a -> IntMap a
    IM.insert Id
    i Id
    bd)

    -- | Increment the fork depth for variables graphed by this action.
    incForkDepthFor :: Grapher a -> Grapher a
    incForkDepthFor :: forall a. Grapher a -> Grapher a
    incForkDepthFor =
      (Env -> Env) -> Grapher a -> Grapher a
    forall a. (Env -> Env) -> Grapher a -> Grapher a
    local ((Env -> Env) -> Grapher a -> Grapher a)
    -> (Env -> Env) -> Grapher a -> Grapher a
    forall a b. (a -> b) -> a -> b
    $ \Env
    env ->
        let meta :: Meta
    meta = Env -> Meta
    envMeta Env
    env
            fork_depth :: Id
    fork_depth = Meta -> Id
    metaForkDepth Meta
    meta
         in Env
    env {envMeta = meta {metaForkDepth = fork_depth + 1}}

    -- | Increment the body depth for variables graphed by this action.
    incBodyDepthFor :: Grapher a -> Grapher a
    incBodyDepthFor :: forall a. Grapher a -> Grapher a
    incBodyDepthFor =
      (Env -> Env) -> Grapher a -> Grapher a
    forall a. (Env -> Env) -> Grapher a -> Grapher a
    local ((Env -> Env) -> Grapher a -> Grapher a)
    -> (Env -> Env) -> Grapher a -> Grapher a
    forall a b. (a -> b) -> a -> b
    $ \Env
    env ->
        let meta :: Meta
    meta = Env -> Meta
    envMeta Env
    env
            body_depth :: Id
    body_depth = Meta -> Id
    metaBodyDepth Meta
    meta
         in Env
    env {envMeta = meta {metaBodyDepth = body_depth + 1}}

    -- | Change the graph id for variables graphed by this action.
    graphIdFor :: Id -> Grapher a -> Grapher a
    graphIdFor :: forall a. Id -> Grapher a -> Grapher a
    graphIdFor Id
    i =
      (Env -> Env) -> Grapher a -> Grapher a
    forall a. (Env -> Env) -> Grapher a -> Grapher a
    local ((Env -> Env) -> Grapher a -> Grapher a)
    -> (Env -> Env) -> Grapher a -> Grapher a
    forall a b. (a -> b) -> a -> b
    $ \Env
    env ->
        let meta :: Meta
    meta = Env -> Meta
    envMeta Env
    env
         in Env
    env {envMeta = meta {metaGraphId = Just i}}

    -- | Capture body stats produced by the given action.
    captureBodyStats :: Grapher a -> Grapher BodyStats
    captureBodyStats :: forall a. Grapher a -> Grapher BodyStats
    captureBodyStats Grapher a
    m = do
      BodyStats
    stats <- (State -> BodyStats) -> Grapher BodyStats
    forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
    gets State -> BodyStats
    stateStats
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateStats = mempty}

      a
    _ <- Grapher a
    m

      BodyStats
    stats' <- (State -> BodyStats) -> Grapher BodyStats
    forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
    gets State -> BodyStats
    stateStats
      (State -> State) -> Grapher ()
    forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
    modify ((State -> State) -> Grapher ()) -> (State -> State) -> Grapher ()
    forall a b. (a -> b) -> a -> b
    $ \State
    st -> State
    st {stateStats = stats <> stats'}

      BodyStats -> Grapher BodyStats
    forall a. a -> StateT State (Reader Env) a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure BodyStats
    stats'

    -- | Can applications of this function be moved to device?
    isHostOnlyFun :: Name -> Grapher Bool
    isHostOnlyFun :: Name -> Grapher Bool
    isHostOnlyFun Name
    fn = (Env -> Bool) -> Grapher Bool
    forall a. (Env -> a) -> Grapher a
    asks ((Env -> Bool) -> Grapher Bool) -> (Env -> Bool) -> Grapher Bool
    forall a b. (a -> b) -> a -> b
    $ Name -> HostOnlyFuns -> Bool
    forall a. Ord a => a -> Set a -> Bool
    S.member Name
    fn (HostOnlyFuns -> Bool) -> (Env -> HostOnlyFuns) -> Env -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Env -> HostOnlyFuns
    envHostOnlyFuns

    -- | Get the 'Meta' corresponding to the current body.
    getMeta :: Grapher Meta
    getMeta :: Grapher Meta
    getMeta = (Env -> Meta) -> Grapher Meta
    forall a. (Env -> a) -> Grapher a
    asks Env -> Meta
    envMeta

    -- | Get the body depth of the current body (its nesting level).
    getBodyDepth :: Grapher BodyDepth
    getBodyDepth :: StateT State (Reader Env) Id
    getBodyDepth = (Env -> Id) -> StateT State (Reader Env) Id
    forall a. (Env -> a) -> Grapher a
    asks (Meta -> Id
    metaBodyDepth (Meta -> Id) -> (Env -> Meta) -> Env -> Id
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Env -> Meta
    envMeta)