    {-# LANGUAGE TypeFamilies #-}

    -- | Interference analysis for Futhark programs.
    module Futhark.Analysis.Interference (Graph, analyseProgGPU) where

    import Control.Monad
    import Control.Monad.Reader
    import Data.Foldable (toList)
    import Data.Function ((&))
    import Data.Functor ((<&>))
    import Data.Map (Map)
    import Data.Map qualified as M
    import Data.Maybe (catMaybes, fromMaybe, mapMaybe)
    import Data.Set (Set)
    import Data.Set qualified as S
    import Futhark.Analysis.Alias qualified as AnlAls
    import Futhark.Analysis.LastUse (LUTabFun)
    import Futhark.Analysis.LastUse qualified as LastUse
    import Futhark.Analysis.MemAlias qualified as MemAlias
    import Futhark.IR.GPUMem
    import Futhark.Util (cartesian, invertMap)

    -- | The set of 'VName' currently in use.
    type InUse = Names

    -- | The set of 'VName' that are no longer in use.
    type LastUsed = Names

    -- | An interference graph. An element @(x, y)@ in the set means that there is
    -- an undirected edge between @x@ and @y@, and therefore the lifetimes of @x@
    -- and @y@ overlap and they "interfere" with each other. We assume that pairs
    -- are always normalized, such that @x@ < @y@, before inserting. This should
    -- prevent any duplicates. We also don't allow any pairs where @x == y@.
    type Graph a = Set (a, a)

    -- | Insert an edge between two values into the graph.
    makeEdge :: (Ord a) => a -> a -> Graph a
    makeEdge :: forall a. Ord a => a -> a -> Graph a
    makeEdge a
    v1 a
    v2
      | a
    v1 a -> a -> Bool
    forall a. Eq a => a -> a -> Bool
    == a
    v2 = Graph a
    forall a. Monoid a => a
    mempty
      | Bool
    otherwise = (a, a) -> Graph a
    forall a. a -> Set a
    S.singleton (a -> a -> a
    forall a. Ord a => a -> a -> a
    min a
    v1 a
    v2, a -> a -> a
    forall a. Ord a => a -> a -> a
    max a
    v1 a
    v2)

    analyseStm ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      Stm GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseStm :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Stm GPUMem -> m (Names, Names, Graph VName)
    analyseStm LUTabFun
    lumap Names
    inuse0 Stm GPUMem
    stm =
      Stm GPUMem
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stm GPUMem
    stm (m (Names, Names, Graph VName) -> m (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $ do
        let pat_name :: VName
    pat_name = PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName (PatElem LetDecMem -> VName) -> PatElem LetDecMem -> VName
    forall a b. (a -> b) -> a -> b
    $ [PatElem LetDecMem] -> PatElem LetDecMem
    forall a. HasCallStack => [a] -> a
    head ([PatElem LetDecMem] -> PatElem LetDecMem)
    -> [PatElem LetDecMem] -> PatElem LetDecMem
    forall a b. (a -> b) -> a -> b
    $ Pat LetDecMem -> [PatElem LetDecMem]
    forall dec. Pat dec -> [PatElem dec]
    patElems (Pat LetDecMem -> [PatElem LetDecMem])
    -> Pat LetDecMem -> [PatElem LetDecMem]
    forall a b. (a -> b) -> a -> b
    $ Stm GPUMem -> Pat (LetDec GPUMem)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPUMem
    stm

        Names
    new_mems <-
          Stm GPUMem -> Pat (LetDec GPUMem)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm GPUMem
    stm
            Pat LetDecMem
    -> (Pat LetDecMem -> [PatElem LetDecMem]) -> [PatElem LetDecMem]
    forall a b. a -> (a -> b) -> b
    & Pat LetDecMem -> [PatElem LetDecMem]
    forall dec. Pat dec -> [PatElem dec]
    patElems
            [PatElem LetDecMem]
    -> ([PatElem LetDecMem] -> m [Maybe VName]) -> m [Maybe VName]
    forall a b. a -> (a -> b) -> b
    & (PatElem LetDecMem -> m (Maybe VName))
    -> [PatElem LetDecMem] -> m [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 (VName -> m (Maybe VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    VName -> m (Maybe VName)
    memInfo (VName -> m (Maybe VName))
    -> (PatElem LetDecMem -> VName)
    -> PatElem LetDecMem
    -> m (Maybe VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PatElem LetDecMem -> VName
    forall dec. PatElem dec -> VName
    patElemName)
            m [Maybe VName] -> ([Maybe VName] -> [VName]) -> m [VName]
    forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
    <&> [Maybe VName] -> [VName]
    forall a. [Maybe a] -> [a]
    catMaybes
            m [VName] -> ([VName] -> Names) -> m Names
    forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
    <&> [VName] -> Names
    namesFromList

        -- `new_mems` should interfere with any mems inside the statement expression
        let inuse_outside :: Names
    inuse_outside = Names
    inuse0 Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    new_mems

        -- `inuse` is the set of memory blocks that are inuse at the end of any code
        -- bodies inside the expression. `lus` is the set of all memory blocks that
        -- have reached their last use in any code bodies inside the
        -- expression. `graph` is the interference graph computed for any code
        -- bodies inside the expression.
        (Names
    inuse, Names
    lus, Graph VName
    graph) <- LUTabFun -> Names -> Exp GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Exp GPUMem -> m (Names, Names, Graph VName)
    analyseExp LUTabFun
    lumap Names
    inuse_outside (Stm GPUMem -> Exp GPUMem
    forall rep. Stm rep -> Exp rep
    stmExp Stm GPUMem
    stm)

        Names
    last_use_mems <-
          VName -> LUTabFun -> Maybe Names
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    pat_name LUTabFun
    lumap
            Maybe Names -> (Maybe Names -> Names) -> Names
    forall a b. a -> (a -> b) -> b
    & Names -> Maybe Names -> Names
    forall a. a -> Maybe a -> a
    fromMaybe Names
    forall a. Monoid a => a
    mempty
            Names -> (Names -> [VName]) -> [VName]
    forall a b. a -> (a -> b) -> b
    & Names -> [VName]
    namesToList
            [VName] -> ([VName] -> m [Maybe VName]) -> m [Maybe VName]
    forall a b. a -> (a -> b) -> b
    & (VName -> m (Maybe VName)) -> [VName] -> m [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 VName -> m (Maybe VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    VName -> m (Maybe VName)
    memInfo
            m [Maybe VName] -> ([Maybe VName] -> [VName]) -> m [VName]
    forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
    <&> [Maybe VName] -> [VName]
    forall a. [Maybe a] -> [a]
    catMaybes
            m [VName] -> ([VName] -> Names) -> m Names
    forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
    <&> [VName] -> Names
    namesFromList
            m Names -> (Names -> Names) -> m Names
    forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
    <&> Names -> Names -> Names
    namesIntersection Names
    inuse_outside

        (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
          ( (Names
    inuse_outside Names -> Names -> Names
    `namesSubtract` Names
    last_use_mems Names -> Names -> Names
    `namesSubtract` Names
    lus)
              Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    new_mems,
            (Names
    lus Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    last_use_mems) Names -> Names -> Names
    `namesSubtract` Names
    new_mems,
            Graph VName
    graph
              Graph VName -> Graph VName -> Graph VName
    forall a. Semigroup a => a -> a -> a
    <> (VName -> VName -> Graph VName)
    -> [VName] -> [VName] -> Graph VName
    forall m (t :: * -> *) a.
    (Monoid m, Foldable t) =>
    (a -> a -> m) -> t a -> t a -> m
    cartesian
                VName -> VName -> Graph VName
    forall a. Ord a => a -> a -> Graph a
    makeEdge
                (Names -> [VName]
    namesToList Names
    inuse_outside)
                (Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Names
    inuse_outside Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    inuse Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    lus Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    last_use_mems)
          )

    -- We conservatively treat all memory arguments to a Loop to
    -- interfere with each other, as well as anything used inside the
    -- loop.  This could potentially be improved by looking at the
    -- interference computed by the loop body wrt. the loop arguments, but
    -- probably very few programs would benefit from this.
    analyseLoopParams ::
      [(FParam GPUMem, SubExp)] ->
      (InUse, LastUsed, Graph VName) ->
      (InUse, LastUsed, Graph VName)
    analyseLoopParams :: [(FParam GPUMem, SubExp)]
    -> (Names, Names, Graph VName) -> (Names, Names, Graph VName)
    analyseLoopParams [(FParam GPUMem, SubExp)]
    merge (Names
    inuse, Names
    lastused, Graph VName
    graph) =
      (Names
    inuse, Names
    lastused, (VName -> VName -> Graph VName)
    -> [VName] -> [VName] -> Graph VName
    forall m (t :: * -> *) a.
    (Monoid m, Foldable t) =>
    (a -> a -> m) -> t a -> t a -> m
    cartesian VName -> VName -> Graph VName
    forall a. Ord a => a -> a -> Graph a
    makeEdge [VName]
    mems ([VName]
    mems [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> [VName]
    inner_mems) Graph VName -> Graph VName -> Graph VName
    forall a. Semigroup a => a -> a -> a
    <> Graph VName
    graph)
      where
        mems :: [VName]
    mems = ((Param (MemInfo SubExp Uniqueness MemBind), SubExp)
     -> Maybe VName)
    -> [(Param (MemInfo SubExp Uniqueness MemBind), SubExp)] -> [VName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe (Param (MemInfo SubExp Uniqueness MemBind), SubExp) -> Maybe VName
    forall {d} {u} {ret}.
    (Param (MemInfo d u ret), SubExp) -> Maybe VName
    isMemArg [(FParam GPUMem, SubExp)]
    [(Param (MemInfo SubExp Uniqueness MemBind), SubExp)]
    merge
        inner_mems :: [VName]
    inner_mems = Names -> [VName]
    namesToList Names
    lastused [VName] -> [VName] -> [VName]
    forall a. Semigroup a => a -> a -> a
    <> Names -> [VName]
    namesToList Names
    inuse
        isMemArg :: (Param (MemInfo d u ret), SubExp) -> Maybe VName
    isMemArg (Param Attrs
    _ VName
    _ MemMem {}, Var VName
    v) = VName -> Maybe VName
    forall a. a -> Maybe a
    Just VName
    v
        isMemArg (Param (MemInfo d u ret), SubExp)
    _ = Maybe VName
    forall a. Maybe a
    Nothing

    analyseExp ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      Exp GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseExp :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Exp GPUMem -> m (Names, Names, Graph VName)
    analyseExp LUTabFun
    lumap Names
    inuse_outside Exp GPUMem
    expr =
      case Exp GPUMem
    expr of
        Match [SubExp]
    _ [Case (Body GPUMem)]
    cases Body GPUMem
    defbody MatchDec (BranchType GPUMem)
    _ ->
          ([(Names, Names, Graph VName)] -> (Names, Names, Graph VName))
    -> m [(Names, Names, Graph VName)] -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [(Names, Names, Graph VName)] -> (Names, Names, Graph VName)
    forall a. Monoid a => [a] -> a
    mconcat (m [(Names, Names, Graph VName)] -> m (Names, Names, Graph VName))
    -> m [(Names, Names, Graph VName)] -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $
            (Body GPUMem -> m (Names, Names, Graph VName))
    -> [Body GPUMem] -> m [(Names, Names, Graph 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 (LUTabFun -> Names -> Body GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Body GPUMem -> m (Names, Names, Graph VName)
    analyseBody LUTabFun
    lumap Names
    inuse_outside) ([Body GPUMem] -> m [(Names, Names, Graph VName)])
    -> [Body GPUMem] -> m [(Names, Names, Graph VName)]
    forall a b. (a -> b) -> a -> b
    $
              Body GPUMem
    defbody Body GPUMem -> [Body GPUMem] -> [Body GPUMem]
    forall a. a -> [a] -> [a]
    : (Case (Body GPUMem) -> Body GPUMem)
    -> [Case (Body GPUMem)] -> [Body GPUMem]
    forall a b. (a -> b) -> [a] -> [b]
    map Case (Body GPUMem) -> Body GPUMem
    forall body. Case body -> body
    caseBody [Case (Body GPUMem)]
    cases
        Loop [(FParam GPUMem, SubExp)]
    merge LoopForm
    _ Body GPUMem
    body ->
          [(FParam GPUMem, SubExp)]
    -> (Names, Names, Graph VName) -> (Names, Names, Graph VName)
    analyseLoopParams [(FParam GPUMem, SubExp)]
    merge ((Names, Names, Graph VName) -> (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> LUTabFun -> Names -> Body GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Body GPUMem -> m (Names, Names, Graph VName)
    analyseBody LUTabFun
    lumap Names
    inuse_outside Body GPUMem
    body
        Op (Inner (SegOp SegOp SegLevel GPUMem
    segop)) -> do
          LUTabFun
    -> Names -> SegOp SegLevel GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *) lvl.
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> SegOp lvl GPUMem -> m (Names, Names, Graph VName)
    analyseSegOp LUTabFun
    lumap Names
    inuse_outside SegOp SegLevel GPUMem
    segop
        Exp GPUMem
    _ ->
          (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Names, Names, Graph VName)
    forall a. Monoid a => a
    mempty

    analyseKernelBody ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      KernelBody GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseKernelBody :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> KernelBody GPUMem -> m (Names, Names, Graph VName)
    analyseKernelBody LUTabFun
    lumap Names
    inuse KernelBody GPUMem
    body = LUTabFun -> Names -> Stms GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Stms GPUMem -> m (Names, Names, Graph VName)
    analyseStms LUTabFun
    lumap Names
    inuse (Stms GPUMem -> m (Names, Names, Graph VName))
    -> Stms GPUMem -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $ KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody GPUMem
    body

    analyseBody ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      Body GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseBody :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Body GPUMem -> m (Names, Names, Graph VName)
    analyseBody LUTabFun
    lumap Names
    inuse Body GPUMem
    body = LUTabFun -> Names -> Stms GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Stms GPUMem -> m (Names, Names, Graph VName)
    analyseStms LUTabFun
    lumap Names
    inuse (Stms GPUMem -> m (Names, Names, Graph VName))
    -> Stms GPUMem -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $ Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms Body GPUMem
    body

    analyseStms ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      Stms GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseStms :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Stms GPUMem -> m (Names, Names, Graph VName)
    analyseStms LUTabFun
    lumap Names
    inuse0 Stms GPUMem
    stms = do
      Stms GPUMem
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stms GPUMem
    stms (m (Names, Names, Graph VName) -> m (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $ ((Names, Names, Graph VName)
     -> Stm GPUMem -> m (Names, Names, Graph VName))
    -> (Names, Names, Graph VName)
    -> [Stm GPUMem]
    -> m (Names, Names, Graph VName)
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM (Names, Names, Graph VName)
    -> Stm GPUMem -> m (Names, Names, Graph VName)
    helper (Names
    inuse0, Names
    forall a. Monoid a => a
    mempty, Graph VName
    forall a. Monoid a => a
    mempty) ([Stm GPUMem] -> m (Names, Names, Graph VName))
    -> [Stm GPUMem] -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $ Stms GPUMem -> [Stm GPUMem]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPUMem
    stms
      where
        helper :: (Names, Names, Graph VName)
    -> Stm GPUMem -> m (Names, Names, Graph VName)
    helper (Names
    inuse, Names
    lus, Graph VName
    graph) Stm GPUMem
    stm = do
          (Names
    inuse', Names
    lus', Graph VName
    graph') <- LUTabFun -> Names -> Stm GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Stm GPUMem -> m (Names, Names, Graph VName)
    analyseStm LUTabFun
    lumap Names
    inuse Stm GPUMem
    stm
          (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Names
    inuse', Names
    lus' Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    lus, Graph VName
    graph' Graph VName -> Graph VName -> Graph VName
    forall a. Semigroup a => a -> a -> a
    <> Graph VName
    graph)

    analyseSegOp ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      SegOp lvl GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseSegOp :: forall (m :: * -> *) lvl.
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> SegOp lvl GPUMem -> m (Names, Names, Graph VName)
    analyseSegOp LUTabFun
    lumap Names
    inuse (SegMap lvl
    _ SegSpace
    _ [Type]
    _ KernelBody GPUMem
    body) =
      LUTabFun
    -> Names -> KernelBody GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> KernelBody GPUMem -> m (Names, Names, Graph VName)
    analyseKernelBody LUTabFun
    lumap Names
    inuse KernelBody GPUMem
    body
    analyseSegOp LUTabFun
    lumap Names
    inuse (SegRed lvl
    _ SegSpace
    _ [Type]
    _ KernelBody GPUMem
    body [SegBinOp GPUMem]
    binops) =
      LUTabFun
    -> Names
    -> [SegBinOp GPUMem]
    -> KernelBody GPUMem
    -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names
    -> [SegBinOp GPUMem]
    -> KernelBody GPUMem
    -> m (Names, Names, Graph VName)
    segWithBinOps LUTabFun
    lumap Names
    inuse [SegBinOp GPUMem]
    binops KernelBody GPUMem
    body
    analyseSegOp LUTabFun
    lumap Names
    inuse (SegScan lvl
    _ SegSpace
    _ [Type]
    _ KernelBody GPUMem
    body [SegBinOp GPUMem]
    binops) = do
      LUTabFun
    -> Names
    -> [SegBinOp GPUMem]
    -> KernelBody GPUMem
    -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names
    -> [SegBinOp GPUMem]
    -> KernelBody GPUMem
    -> m (Names, Names, Graph VName)
    segWithBinOps LUTabFun
    lumap Names
    inuse [SegBinOp GPUMem]
    binops KernelBody GPUMem
    body
    analyseSegOp LUTabFun
    lumap Names
    inuse (SegHist lvl
    _ SegSpace
    _ [Type]
    _ KernelBody GPUMem
    body [HistOp GPUMem]
    histops) = do
      (Names
    inuse', Names
    lus', Graph VName
    graph) <- LUTabFun
    -> Names -> KernelBody GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> KernelBody GPUMem -> m (Names, Names, Graph VName)
    analyseKernelBody LUTabFun
    lumap Names
    inuse KernelBody GPUMem
    body
      (Names
    inuse'', Names
    lus'', Graph VName
    graph') <- [(Names, Names, Graph VName)] -> (Names, Names, Graph VName)
    forall a. Monoid a => [a] -> a
    mconcat ([(Names, Names, Graph VName)] -> (Names, Names, Graph VName))
    -> m [(Names, Names, Graph VName)] -> m (Names, Names, Graph VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (HistOp GPUMem -> m (Names, Names, Graph VName))
    -> [HistOp GPUMem] -> m [(Names, Names, Graph 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 (LUTabFun -> Names -> HistOp GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> HistOp GPUMem -> m (Names, Names, Graph VName)
    analyseHistOp LUTabFun
    lumap Names
    inuse') [HistOp GPUMem]
    histops
      (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Names
    inuse'', Names
    lus' Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    lus'', Graph VName
    graph Graph VName -> Graph VName -> Graph VName
    forall a. Semigroup a => a -> a -> a
    <> Graph VName
    graph')

    segWithBinOps ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      [SegBinOp GPUMem] ->
      KernelBody GPUMem ->
      m (InUse, LastUsed, Graph VName)
    segWithBinOps :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names
    -> [SegBinOp GPUMem]
    -> KernelBody GPUMem
    -> m (Names, Names, Graph VName)
    segWithBinOps LUTabFun
    lumap Names
    inuse [SegBinOp GPUMem]
    binops KernelBody GPUMem
    body = do
      (Names
    inuse', Names
    lus', Graph VName
    graph) <- LUTabFun
    -> Names -> KernelBody GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> KernelBody GPUMem -> m (Names, Names, Graph VName)
    analyseKernelBody LUTabFun
    lumap Names
    inuse KernelBody GPUMem
    body
      (Names
    inuse'', Names
    lus'', Graph VName
    graph') <-
        [(Names, Names, Graph VName)] -> (Names, Names, Graph VName)
    forall a. Monoid a => [a] -> a
    mconcat
          ([(Names, Names, Graph VName)] -> (Names, Names, Graph VName))
    -> m [(Names, Names, Graph VName)] -> m (Names, Names, Graph VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SegBinOp GPUMem -> m (Names, Names, Graph VName))
    -> [SegBinOp GPUMem] -> m [(Names, Names, Graph 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
            (LUTabFun
    -> Names -> SegBinOp GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> SegBinOp GPUMem -> m (Names, Names, Graph VName)
    analyseSegBinOp LUTabFun
    lumap Names
    inuse')
            [SegBinOp GPUMem]
    binops
      (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Names
    inuse'', Names
    lus' Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    lus'', Graph VName
    graph Graph VName -> Graph VName -> Graph VName
    forall a. Semigroup a => a -> a -> a
    <> Graph VName
    graph')

    analyseSegBinOp ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      SegBinOp GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseSegBinOp :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> SegBinOp GPUMem -> m (Names, Names, Graph VName)
    analyseSegBinOp LUTabFun
    lumap Names
    inuse (SegBinOp Commutativity
    _ Lambda GPUMem
    lambda [SubExp]
    _ Shape
    _) =
      LUTabFun -> Names -> Lambda GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Lambda GPUMem -> m (Names, Names, Graph VName)
    analyseLambda LUTabFun
    lumap Names
    inuse Lambda GPUMem
    lambda

    analyseHistOp ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      HistOp GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseHistOp :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> HistOp GPUMem -> m (Names, Names, Graph VName)
    analyseHistOp LUTabFun
    lumap Names
    inuse = LUTabFun -> Names -> Lambda GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Lambda GPUMem -> m (Names, Names, Graph VName)
    analyseLambda LUTabFun
    lumap Names
    inuse (Lambda GPUMem -> m (Names, Names, Graph VName))
    -> (HistOp GPUMem -> Lambda GPUMem)
    -> HistOp GPUMem
    -> m (Names, Names, Graph VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . HistOp GPUMem -> Lambda GPUMem
    forall rep. HistOp rep -> Lambda rep
    histOp

    analyseLambda ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      InUse ->
      Lambda GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseLambda :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Lambda GPUMem -> m (Names, Names, Graph VName)
    analyseLambda LUTabFun
    lumap Names
    inuse = LUTabFun -> Names -> Body GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Names -> Body GPUMem -> m (Names, Names, Graph VName)
    analyseBody LUTabFun
    lumap Names
    inuse (Body GPUMem -> m (Names, Names, Graph VName))
    -> (Lambda GPUMem -> Body GPUMem)
    -> Lambda GPUMem
    -> m (Names, Names, Graph VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Lambda GPUMem -> Body GPUMem
    forall rep. Lambda rep -> Body rep
    lambdaBody

    analyseProgGPU :: Prog GPUMem -> Graph VName
    analyseProgGPU :: Prog GPUMem -> Graph VName
    analyseProgGPU Prog GPUMem
    prog = Stms GPUMem -> Graph VName
    onConsts (Prog GPUMem -> Stms GPUMem
    forall rep. Prog rep -> Stms rep
    progConsts Prog GPUMem
    prog) Graph VName -> Graph VName -> Graph VName
    forall a. Semigroup a => a -> a -> a
    <> (FunDef GPUMem -> Graph VName) -> [FunDef GPUMem] -> Graph VName
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap FunDef GPUMem -> Graph VName
    onFun (Prog GPUMem -> [FunDef GPUMem]
    forall rep. Prog rep -> [FunDef rep]
    progFuns Prog GPUMem
    prog)
      where
        (MemAliases
    consts_aliases, Map Name MemAliases
    funs_aliases) = Prog GPUMem -> (MemAliases, Map Name MemAliases)
    MemAlias.analyzeGPUMem Prog GPUMem
    prog
        (LUTabFun
    lumap_consts, Map Name LUTabFun
    lumap) = Prog (Aliases GPUMem) -> (LUTabFun, Map Name LUTabFun)
    LastUse.lastUseGPUMem (Prog (Aliases GPUMem) -> (LUTabFun, Map Name LUTabFun))
    -> Prog (Aliases GPUMem) -> (LUTabFun, Map Name LUTabFun)
    forall a b. (a -> b) -> a -> b
    $ Prog GPUMem -> Prog (Aliases GPUMem)
    forall rep. AliasableRep rep => Prog rep -> Prog (Aliases rep)
    AnlAls.aliasAnalysis Prog GPUMem
    prog
        onFun :: FunDef GPUMem -> Graph VName
    onFun FunDef GPUMem
    f =
          MemAliases -> Graph VName -> Graph VName
    applyAliases (MemAliases -> Maybe MemAliases -> MemAliases
    forall a. a -> Maybe a -> a
    fromMaybe MemAliases
    forall a. Monoid a => a
    mempty (Maybe MemAliases -> MemAliases) -> Maybe MemAliases -> MemAliases
    forall a b. (a -> b) -> a -> b
    $ Name -> Map Name MemAliases -> Maybe MemAliases
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup (FunDef GPUMem -> Name
    forall rep. FunDef rep -> Name
    funDefName FunDef GPUMem
    f) Map Name MemAliases
    funs_aliases) (Graph VName -> Graph VName) -> Graph VName -> Graph VName
    forall a b. (a -> b) -> a -> b
    $
            Reader (Scope GPUMem) (Graph VName) -> Scope GPUMem -> Graph VName
    forall r a. Reader r a -> r -> a
    runReader (LUTabFun -> Stms GPUMem -> Reader (Scope GPUMem) (Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Stms GPUMem -> m (Graph VName)
    analyseGPU (Map Name LUTabFun
    lumap Map Name LUTabFun -> Name -> LUTabFun
    forall k a. Ord k => Map k a -> k -> a
    M.! FunDef GPUMem -> Name
    forall rep. FunDef rep -> Name
    funDefName FunDef GPUMem
    f) (Stms GPUMem -> Reader (Scope GPUMem) (Graph VName))
    -> Stms GPUMem -> Reader (Scope GPUMem) (Graph VName)
    forall a b. (a -> b) -> a -> b
    $ Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms (Body GPUMem -> Stms GPUMem) -> Body GPUMem -> Stms GPUMem
    forall a b. (a -> b) -> a -> b
    $ FunDef GPUMem -> Body GPUMem
    forall rep. FunDef rep -> Body rep
    funDefBody FunDef GPUMem
    f) (Scope GPUMem -> Graph VName) -> Scope GPUMem -> Graph VName
    forall a b. (a -> b) -> a -> b
    $
              FunDef GPUMem -> Scope GPUMem
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf FunDef GPUMem
    f
        onConsts :: Stms GPUMem -> Graph VName
    onConsts Stms GPUMem
    stms =
          MemAliases -> Graph VName -> Graph VName
    applyAliases MemAliases
    consts_aliases (Graph VName -> Graph VName) -> Graph VName -> Graph VName
    forall a b. (a -> b) -> a -> b
    $
            Reader (Scope GPUMem) (Graph VName) -> Scope GPUMem -> Graph VName
    forall r a. Reader r a -> r -> a
    runReader (LUTabFun -> Stms GPUMem -> Reader (Scope GPUMem) (Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Stms GPUMem -> m (Graph VName)
    analyseGPU LUTabFun
    lumap_consts Stms GPUMem
    stms) (Scope GPUMem
    forall a. Monoid a => a
    mempty :: Scope GPUMem)

    applyAliases :: MemAlias.MemAliases -> Graph VName -> Graph VName
    applyAliases :: MemAliases -> Graph VName -> Graph VName
    applyAliases MemAliases
    aliases =
      -- For each pair @(x, y)@ in graph, all memory aliases of x should interfere with all memory aliases of y
      ((VName, VName) -> Graph VName) -> Graph VName -> Graph VName
    forall m a. Monoid m => (a -> m) -> Set a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap
        ( \(VName
    x, VName
    y) ->
            let xs :: Names
    xs = MemAliases -> VName -> Names
    MemAlias.aliasesOf MemAliases
    aliases VName
    x Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> VName -> Names
    oneName VName
    x
                ys :: Names
    ys = MemAliases -> VName -> Names
    MemAlias.aliasesOf MemAliases
    aliases VName
    y Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> VName -> Names
    oneName VName
    y
             in (VName -> VName -> Graph VName)
    -> [VName] -> [VName] -> Graph VName
    forall m (t :: * -> *) a.
    (Monoid m, Foldable t) =>
    (a -> a -> m) -> t a -> t a -> m
    cartesian VName -> VName -> Graph VName
    forall a. Ord a => a -> a -> Graph a
    makeEdge (Names -> [VName]
    namesToList Names
    xs) (Names -> [VName]
    namesToList Names
    ys)
        )

    -- | Perform interference analysis on the given statements. The result is a
    -- triple of the names currently in use, names that hit their last use somewhere
    -- within, and the resulting graph.
    analyseGPU ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      Stms GPUMem ->
      m (Graph VName)
    analyseGPU :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Stms GPUMem -> m (Graph VName)
    analyseGPU LUTabFun
    lumap Stms GPUMem
    stms = do
      (Names
    _, Names
    _, Graph VName
    graph) <- LUTabFun -> Stms GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Stms GPUMem -> m (Names, Names, Graph VName)
    analyseGPU' LUTabFun
    lumap Stms GPUMem
    stms
      -- We need to insert edges between memory blocks which differ in size, if they
      -- are in DefaultSpace. The problem is that during memory expansion,
      -- DefaultSpace arrays in kernels are interleaved. If the element sizes of two
      -- merged memory blocks are different, threads might try to read and write to
      -- overlapping memory positions. More information here:
      -- https://munksgaard.me/technical-diary/2020-12-30.html#org210775b
      Map VName Space
    spaces <- (Space -> Bool) -> Map VName Space -> Map VName Space
    forall a k. (a -> Bool) -> Map k a -> Map k a
    M.filter (Space -> Space -> Bool
    forall a. Eq a => a -> a -> Bool
    == Space
    DefaultSpace) (Map VName Space -> Map VName Space)
    -> m (Map VName Space) -> m (Map VName Space)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Stms GPUMem -> m (Map VName Space)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stms GPUMem -> m (Map VName Space)
    memSpaces Stms GPUMem
    stms
      Map Int (Set VName)
    inv_size_map <-
        Stms GPUMem -> m (Map VName Int)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stms GPUMem -> m (Map VName Int)
    memSizes Stms GPUMem
    stms
          m (Map VName Int)
    -> (Map VName Int -> Map VName Int) -> m (Map VName Int)
    forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
    <&> (Map VName Int -> Set VName -> Map VName Int)
    -> Set VName -> Map VName Int -> Map VName Int
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Map VName Int -> Set VName -> Map VName Int
    forall k a. Ord k => Map k a -> Set k -> Map k a
    M.restrictKeys ([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
    $ Map VName Space -> [VName]
    forall k a. Map k a -> [k]
    M.keys Map VName Space
    spaces)
          m (Map VName Int)
    -> (Map VName Int -> Map Int (Set VName))
    -> m (Map Int (Set VName))
    forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
    <&> Map VName Int -> Map Int (Set VName)
    forall v k. (Ord v, Ord k) => Map k v -> Map v (Set k)
    invertMap
      let new_edges :: Graph VName
    new_edges =
            (Set VName -> Set VName -> Graph VName)
    -> Map Int (Set VName) -> Map Int (Set VName) -> Graph VName
    forall m (t :: * -> *) a.
    (Monoid m, Foldable t) =>
    (a -> a -> m) -> t a -> t a -> m
    cartesian
              (\Set VName
    x Set VName
    y -> if Set VName
    x Set VName -> Set VName -> Bool
    forall a. Eq a => a -> a -> Bool
    /= Set VName
    y then (VName -> VName -> Graph VName)
    -> Set VName -> Set VName -> Graph VName
    forall m (t :: * -> *) a.
    (Monoid m, Foldable t) =>
    (a -> a -> m) -> t a -> t a -> m
    cartesian VName -> VName -> Graph VName
    forall a. Ord a => a -> a -> Graph a
    makeEdge Set VName
    x Set VName
    y else Graph VName
    forall a. Monoid a => a
    mempty)
              Map Int (Set VName)
    inv_size_map
              Map Int (Set VName)
    inv_size_map
      Graph VName -> m (Graph VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Graph VName -> m (Graph VName)) -> Graph VName -> m (Graph VName)
    forall a b. (a -> b) -> a -> b
    $ Graph VName
    graph Graph VName -> Graph VName -> Graph VName
    forall a. Semigroup a => a -> a -> a
    <> Graph VName
    new_edges

    -- | Return a mapping from memory blocks to their element sizes in the given
    -- statements.
    memSizes :: (LocalScope GPUMem m) => Stms GPUMem -> m (Map VName Int)
    memSizes :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stms GPUMem -> m (Map VName Int)
    memSizes Stms GPUMem
    stms =
      Stms GPUMem -> m (Map VName Int) -> m (Map VName Int)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stms GPUMem
    stms (m (Map VName Int) -> m (Map VName Int))
    -> m (Map VName Int) -> m (Map VName Int)
    forall a b. (a -> b) -> a -> b
    $ ([Map VName Int] -> Map VName Int)
    -> m [Map VName Int] -> m (Map VName Int)
    forall a b. (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [Map VName Int] -> Map VName Int
    forall a. Monoid a => [a] -> a
    mconcat (m [Map VName Int] -> m (Map VName Int))
    -> ([Stm GPUMem] -> m [Map VName Int])
    -> [Stm GPUMem]
    -> m (Map VName Int)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Stm GPUMem -> m (Map VName Int))
    -> [Stm GPUMem] -> m [Map VName Int]
    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 Stm GPUMem -> m (Map VName Int)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stm GPUMem -> m (Map VName Int)
    memSizesStm ([Stm GPUMem] -> m (Map VName Int))
    -> [Stm GPUMem] -> m (Map VName Int)
    forall a b. (a -> b) -> a -> b
    $ Stms GPUMem -> [Stm GPUMem]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms GPUMem
    stms
      where
        memSizesStm :: (LocalScope GPUMem m) => Stm GPUMem -> m (Map VName Int)
        memSizesStm :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stm GPUMem -> m (Map VName Int)
    memSizesStm (Let Pat (LetDec GPUMem)
    pat StmAux (ExpDec GPUMem)
    _ Exp GPUMem
    e) = do
          Map VName Int
    arraySizes <- ([Map VName Int] -> Map VName Int)
    -> m [Map VName Int] -> m (Map VName Int)
    forall a b. (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [Map VName Int] -> Map VName Int
    forall a. Monoid a => [a] -> a
    mconcat (m [Map VName Int] -> m (Map VName Int))
    -> ([VName] -> m [Map VName Int]) -> [VName] -> m (Map VName Int)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (VName -> m (Map VName Int)) -> [VName] -> m [Map VName Int]
    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 -> m (Map VName Int)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    VName -> m (Map VName Int)
    memElemSize ([VName] -> m (Map VName Int)) -> [VName] -> m (Map VName Int)
    forall a b. (a -> b) -> a -> b
    $ Pat LetDecMem -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat (LetDec GPUMem)
    Pat LetDecMem
    pat
          Map VName Int
    arraySizes' <- Exp GPUMem -> m (Map VName Int)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    Exp GPUMem -> m (Map VName Int)
    memSizesExp Exp GPUMem
    e
          Map VName Int -> m (Map VName Int)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Map VName Int -> m (Map VName Int))
    -> Map VName Int -> m (Map VName Int)
    forall a b. (a -> b) -> a -> b
    $ Map VName Int
    arraySizes Map VName Int -> Map VName Int -> Map VName Int
    forall a. Semigroup a => a -> a -> a
    <> Map VName Int
    arraySizes'
        memSizesExp :: (LocalScope GPUMem m) => Exp GPUMem -> m (Map VName Int)
        memSizesExp :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    Exp GPUMem -> m (Map VName Int)
    memSizesExp (Op (Inner (SegOp SegOp SegLevel GPUMem
    segop))) =
          let body :: KernelBody GPUMem
    body = SegOp SegLevel GPUMem -> KernelBody GPUMem
    forall lvl rep. SegOp lvl rep -> KernelBody rep
    segBody SegOp SegLevel GPUMem
    segop
           in Stms GPUMem -> m (Map VName Int) -> m (Map VName Int)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf (KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody GPUMem
    body)
                (m (Map VName Int) -> m (Map VName Int))
    -> m (Map VName Int) -> m (Map VName Int)
    forall a b. (a -> b) -> a -> b
    $ ([Map VName Int] -> Map VName Int)
    -> m [Map VName Int] -> m (Map VName Int)
    forall a b. (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [Map VName Int] -> Map VName Int
    forall a. Monoid a => [a] -> a
    mconcat
                  (m [Map VName Int] -> m (Map VName Int))
    -> ([Stm GPUMem] -> m [Map VName Int])
    -> [Stm GPUMem]
    -> m (Map VName Int)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Stm GPUMem -> m (Map VName Int))
    -> [Stm GPUMem] -> m [Map VName Int]
    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 Stm GPUMem -> m (Map VName Int)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stm GPUMem -> m (Map VName Int)
    memSizesStm
                ([Stm GPUMem] -> m (Map VName Int))
    -> [Stm GPUMem] -> m (Map VName Int)
    forall a b. (a -> b) -> a -> b
    $ Stms GPUMem -> [Stm GPUMem]
    forall rep. Stms rep -> [Stm rep]
    stmsToList
                (Stms GPUMem -> [Stm GPUMem]) -> Stms GPUMem -> [Stm GPUMem]
    forall a b. (a -> b) -> a -> b
    $ KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms KernelBody GPUMem
    body
        memSizesExp (Match [SubExp]
    _ [Case (Body GPUMem)]
    cases Body GPUMem
    defbody MatchDec (BranchType GPUMem)
    _) = do
          [Map VName Int] -> Map VName Int
    forall a. Monoid a => [a] -> a
    mconcat ([Map VName Int] -> Map VName Int)
    -> m [Map VName Int] -> m (Map VName Int)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Body GPUMem -> m (Map VName Int))
    -> [Body GPUMem] -> m [Map VName Int]
    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 (Stms GPUMem -> m (Map VName Int)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stms GPUMem -> m (Map VName Int)
    memSizes (Stms GPUMem -> m (Map VName Int))
    -> (Body GPUMem -> Stms GPUMem) -> Body GPUMem -> m (Map VName Int)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms) (Body GPUMem
    defbody Body GPUMem -> [Body GPUMem] -> [Body GPUMem]
    forall a. a -> [a] -> [a]
    : (Case (Body GPUMem) -> Body GPUMem)
    -> [Case (Body GPUMem)] -> [Body GPUMem]
    forall a b. (a -> b) -> [a] -> [b]
    map Case (Body GPUMem) -> Body GPUMem
    forall body. Case body -> body
    caseBody [Case (Body GPUMem)]
    cases)
        memSizesExp (Loop [(FParam GPUMem, SubExp)]
    _ LoopForm
    _ Body GPUMem
    body) =
          Stms GPUMem -> m (Map VName Int)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stms GPUMem -> m (Map VName Int)
    memSizes (Stms GPUMem -> m (Map VName Int))
    -> Stms GPUMem -> m (Map VName Int)
    forall a b. (a -> b) -> a -> b
    $ Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms Body GPUMem
    body
        memSizesExp Exp GPUMem
    _ = Map VName Int -> m (Map VName Int)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Map VName Int
    forall a. Monoid a => a
    mempty

    -- | Return a mapping from memory blocks to the space they are allocated in.
    memSpaces :: (LocalScope GPUMem m) => Stms GPUMem -> m (Map VName Space)
    memSpaces :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stms GPUMem -> m (Map VName Space)
    memSpaces Stms GPUMem
    stms =
      Map VName Space -> m (Map VName Space)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Map VName Space -> m (Map VName Space))
    -> Map VName Space -> m (Map VName Space)
    forall a b. (a -> b) -> a -> b
    $ (Stm GPUMem -> Map VName Space) -> Stms GPUMem -> Map VName Space
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Stm GPUMem -> Map VName Space
    getSpacesStm Stms GPUMem
    stms
      where
        getSpacesStm :: Stm GPUMem -> Map VName Space
        getSpacesStm :: Stm GPUMem -> Map VName Space
    getSpacesStm (Let (Pat [PatElem VName
    name LetDec GPUMem
    _]) StmAux (ExpDec GPUMem)
    _ (Op (Alloc SubExp
    _ Space
    sp))) =
          VName -> Space -> Map VName Space
    forall k a. k -> a -> Map k a
    M.singleton VName
    name Space
    sp
        getSpacesStm (Let Pat (LetDec GPUMem)
    _ StmAux (ExpDec GPUMem)
    _ (Op (Alloc SubExp
    _ Space
    _))) = [Char] -> Map VName Space
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "impossible"
        getSpacesStm (Let Pat (LetDec GPUMem)
    _ StmAux (ExpDec GPUMem)
    _ (Op (Inner (SegOp SegOp SegLevel GPUMem
    segop)))) =
          (Stm GPUMem -> Map VName Space) -> Stms GPUMem -> Map VName Space
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Stm GPUMem -> Map VName Space
    getSpacesStm (Stms GPUMem -> Map VName Space) -> Stms GPUMem -> Map VName Space
    forall a b. (a -> b) -> a -> b
    $ KernelBody GPUMem -> Stms GPUMem
    forall rep. KernelBody rep -> Stms rep
    kernelBodyStms (KernelBody GPUMem -> Stms GPUMem)
    -> KernelBody GPUMem -> Stms GPUMem
    forall a b. (a -> b) -> a -> b
    $ SegOp SegLevel GPUMem -> KernelBody GPUMem
    forall lvl rep. SegOp lvl rep -> KernelBody rep
    segBody SegOp SegLevel GPUMem
    segop
        getSpacesStm (Let Pat (LetDec GPUMem)
    _ StmAux (ExpDec GPUMem)
    _ (Match [SubExp]
    _ [Case (Body GPUMem)]
    cases Body GPUMem
    defbody MatchDec (BranchType GPUMem)
    _)) =
          (Body GPUMem -> Map VName Space)
    -> [Body GPUMem] -> Map VName Space
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap ((Stm GPUMem -> Map VName Space) -> Stms GPUMem -> Map VName Space
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Stm GPUMem -> Map VName Space
    getSpacesStm (Stms GPUMem -> Map VName Space)
    -> (Body GPUMem -> Stms GPUMem) -> Body GPUMem -> Map VName Space
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms) ([Body GPUMem] -> Map VName Space)
    -> [Body GPUMem] -> Map VName Space
    forall a b. (a -> b) -> a -> b
    $ Body GPUMem
    defbody Body GPUMem -> [Body GPUMem] -> [Body GPUMem]
    forall a. a -> [a] -> [a]
    : (Case (Body GPUMem) -> Body GPUMem)
    -> [Case (Body GPUMem)] -> [Body GPUMem]
    forall a b. (a -> b) -> [a] -> [b]
    map Case (Body GPUMem) -> Body GPUMem
    forall body. Case body -> body
    caseBody [Case (Body GPUMem)]
    cases
        getSpacesStm (Let Pat (LetDec GPUMem)
    _ StmAux (ExpDec GPUMem)
    _ (Loop [(FParam GPUMem, SubExp)]
    _ LoopForm
    _ Body GPUMem
    body)) =
          (Stm GPUMem -> Map VName Space) -> Stms GPUMem -> Map VName Space
    forall m a. Monoid m => (a -> m) -> Seq a -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap Stm GPUMem -> Map VName Space
    getSpacesStm (Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms Body GPUMem
    body)
        getSpacesStm Stm GPUMem
    _ = Map VName Space
    forall a. Monoid a => a
    mempty

    analyseGPU' ::
      (LocalScope GPUMem m) =>
      LUTabFun ->
      Stms GPUMem ->
      m (InUse, LastUsed, Graph VName)
    analyseGPU' :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Stms GPUMem -> m (Names, Names, Graph VName)
    analyseGPU' LUTabFun
    lumap Stms GPUMem
    stms =
      [(Names, Names, Graph VName)] -> (Names, Names, Graph VName)
    forall a. Monoid a => [a] -> a
    mconcat ([(Names, Names, Graph VName)] -> (Names, Names, Graph VName))
    -> (Seq (Names, Names, Graph VName)
        -> [(Names, Names, Graph VName)])
    -> Seq (Names, Names, Graph VName)
    -> (Names, Names, Graph VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Seq (Names, Names, Graph VName) -> [(Names, Names, Graph VName)]
    forall a. Seq a -> [a]
    forall (t :: * -> *) a. Foldable t => t a -> [a]
    toList (Seq (Names, Names, Graph VName) -> (Names, Names, Graph VName))
    -> m (Seq (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Stm GPUMem -> m (Names, Names, Graph VName))
    -> Stms GPUMem -> m (Seq (Names, Names, Graph 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) -> Seq a -> m (Seq b)
    mapM Stm GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stm GPUMem -> m (Names, Names, Graph VName)
    helper Stms GPUMem
    stms
      where
        helper ::
          (LocalScope GPUMem m) =>
          Stm GPUMem ->
          m (InUse, LastUsed, Graph VName)
        helper :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    Stm GPUMem -> m (Names, Names, Graph VName)
    helper stm :: Stm GPUMem
    stm@Let {stmExp :: forall rep. Stm rep -> Exp rep
    stmExp = Op (Inner (SegOp SegOp SegLevel GPUMem
    segop))} =
          Stm GPUMem
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stm GPUMem
    stm (m (Names, Names, Graph VName) -> m (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $ LUTabFun
    -> Names -> SegOp SegLevel GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *) lvl.
    LocalScope GPUMem m =>
    LUTabFun
    -> Names -> SegOp lvl GPUMem -> m (Names, Names, Graph VName)
    analyseSegOp LUTabFun
    lumap Names
    forall a. Monoid a => a
    mempty SegOp SegLevel GPUMem
    segop
        helper stm :: Stm GPUMem
    stm@Let {stmExp :: forall rep. Stm rep -> Exp rep
    stmExp = Match [SubExp]
    _ [Case (Body GPUMem)]
    cases Body GPUMem
    defbody MatchDec (BranchType GPUMem)
    _} =
          Stm GPUMem
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stm GPUMem
    stm (m (Names, Names, Graph VName) -> m (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $
            [(Names, Names, Graph VName)] -> (Names, Names, Graph VName)
    forall a. Monoid a => [a] -> a
    mconcat
              ([(Names, Names, Graph VName)] -> (Names, Names, Graph VName))
    -> m [(Names, Names, Graph VName)] -> m (Names, Names, Graph VName)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (Body GPUMem -> m (Names, Names, Graph VName))
    -> [Body GPUMem] -> m [(Names, Names, Graph 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 (LUTabFun -> Stms GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Stms GPUMem -> m (Names, Names, Graph VName)
    analyseGPU' LUTabFun
    lumap (Stms GPUMem -> m (Names, Names, Graph VName))
    -> (Body GPUMem -> Stms GPUMem)
    -> Body GPUMem
    -> m (Names, Names, Graph VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms) (Body GPUMem
    defbody Body GPUMem -> [Body GPUMem] -> [Body GPUMem]
    forall a. a -> [a] -> [a]
    : (Case (Body GPUMem) -> Body GPUMem)
    -> [Case (Body GPUMem)] -> [Body GPUMem]
    forall a b. (a -> b) -> [a] -> [b]
    map Case (Body GPUMem) -> Body GPUMem
    forall body. Case body -> body
    caseBody [Case (Body GPUMem)]
    cases)
        helper stm :: Stm GPUMem
    stm@Let {stmExp :: forall rep. Stm rep -> Exp rep
    stmExp = Loop [(FParam GPUMem, SubExp)]
    merge LoopForm
    _ Body GPUMem
    body} =
          ((Names, Names, Graph VName) -> (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap ([(FParam GPUMem, SubExp)]
    -> (Names, Names, Graph VName) -> (Names, Names, Graph VName)
    analyseLoopParams [(FParam GPUMem, SubExp)]
    merge) (m (Names, Names, Graph VName) -> m (Names, Names, Graph VName))
    -> (m (Names, Names, Graph VName) -> m (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName)
    -> m (Names, Names, Graph VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Stm GPUMem
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stm GPUMem
    stm (m (Names, Names, Graph VName) -> m (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $
            LUTabFun -> Stms GPUMem -> m (Names, Names, Graph VName)
    forall (m :: * -> *).
    LocalScope GPUMem m =>
    LUTabFun -> Stms GPUMem -> m (Names, Names, Graph VName)
    analyseGPU' LUTabFun
    lumap (Stms GPUMem -> m (Names, Names, Graph VName))
    -> Stms GPUMem -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $
              Body GPUMem -> Stms GPUMem
    forall rep. Body rep -> Stms rep
    bodyStms Body GPUMem
    body
        helper Stm GPUMem
    stm =
          Stm GPUMem
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall rep a (m :: * -> *) b.
    (Scoped rep a, LocalScope rep m) =>
    a -> m b -> m b
    inScopeOf Stm GPUMem
    stm (m (Names, Names, Graph VName) -> m (Names, Names, Graph VName))
    -> m (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a b. (a -> b) -> a -> b
    $ (Names, Names, Graph VName) -> m (Names, Names, Graph VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Names, Names, Graph VName)
    forall a. Monoid a => a
    mempty

    memInfo :: (LocalScope GPUMem m) => VName -> m (Maybe VName)
    memInfo :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    VName -> m (Maybe VName)
    memInfo VName
    vname = do
      Maybe LetDecMem
    summary <- (Scope GPUMem -> Maybe LetDecMem) -> m (Maybe LetDecMem)
    forall a. (Scope GPUMem -> a) -> m a
    forall rep (m :: * -> *) a.
    HasScope rep m =>
    (Scope rep -> a) -> m a
    asksScope ((NameInfo GPUMem -> LetDecMem)
    -> Maybe (NameInfo GPUMem) -> Maybe LetDecMem
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap NameInfo GPUMem -> LetDecMem
    forall rep (inner :: * -> *).
    Mem rep inner =>
    NameInfo rep -> LetDecMem
    nameInfoToMemInfo (Maybe (NameInfo GPUMem) -> Maybe LetDecMem)
    -> (Scope GPUMem -> Maybe (NameInfo GPUMem))
    -> Scope GPUMem
    -> Maybe LetDecMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> Scope GPUMem -> Maybe (NameInfo GPUMem)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    vname)
      case Maybe LetDecMem
    summary of
        Just (MemArray PrimType
    _ Shape
    _ NoUniqueness
    _ (ArrayIn VName
    mem LMAD
    _)) ->
          Maybe VName -> m (Maybe VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Maybe VName -> m (Maybe VName)) -> Maybe VName -> m (Maybe VName)
    forall a b. (a -> b) -> a -> b
    $ VName -> Maybe VName
    forall a. a -> Maybe a
    Just VName
    mem
        Maybe LetDecMem
    _ ->
          Maybe VName -> m (Maybe VName)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe VName
    forall a. Maybe a
    Nothing

    -- | Returns a mapping from memory block to element size. The input is the
    -- `VName` of a variable (supposedly an array), and the result is a mapping from
    -- the memory block of that array to element size of the array.
    memElemSize :: (LocalScope GPUMem m) => VName -> m (Map VName Int)
    memElemSize :: forall (m :: * -> *).
    LocalScope GPUMem m =>
    VName -> m (Map VName Int)
    memElemSize VName
    vname = do
      Maybe LetDecMem
    summary <- (Scope GPUMem -> Maybe LetDecMem) -> m (Maybe LetDecMem)
    forall a. (Scope GPUMem -> a) -> m a
    forall rep (m :: * -> *) a.
    HasScope rep m =>
    (Scope rep -> a) -> m a
    asksScope ((NameInfo GPUMem -> LetDecMem)
    -> Maybe (NameInfo GPUMem) -> Maybe LetDecMem
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap NameInfo GPUMem -> LetDecMem
    forall rep (inner :: * -> *).
    Mem rep inner =>
    NameInfo rep -> LetDecMem
    nameInfoToMemInfo (Maybe (NameInfo GPUMem) -> Maybe LetDecMem)
    -> (Scope GPUMem -> Maybe (NameInfo GPUMem))
    -> Scope GPUMem
    -> Maybe LetDecMem
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VName -> Scope GPUMem -> Maybe (NameInfo GPUMem)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    vname)
      case Maybe LetDecMem
    summary of
        Just (MemArray PrimType
    pt Shape
    _ NoUniqueness
    _ (ArrayIn VName
    mem LMAD
    _)) ->
          Map VName Int -> m (Map VName Int)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Map VName Int -> m (Map VName Int))
    -> Map VName Int -> m (Map VName Int)
    forall a b. (a -> b) -> a -> b
    $ VName -> Int -> Map VName Int
    forall k a. k -> a -> Map k a
    M.singleton VName
    mem (PrimType -> Int
    forall a. Num a => PrimType -> a
    primByteSize PrimType
    pt)
        Maybe LetDecMem
    _ ->
          Map VName Int -> m (Map VName Int)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Map VName Int
    forall a. Monoid a => a
    mempty